import gradio as gr
from main import main
from arguments import parse_args

def generate_image(prompt):
    # Set up arguments
    args = parse_args()
    args.task = "single"
    args.prompt = prompt
    args.model = "sd-turbo"  # or another supported model
    args.save_dir = "./outputs"
    args.save_all_images = True
    
    # Run the main function
    main(args)
    
    # Return the path to the generated image
    return f"{args.save_dir}/{args.task}/{args.model}_{args.prompt}/best_image.png"

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="ReNO Image Generation",
    description="Enter a prompt to generate an image using ReNO."
)

# Launch the app
iface.launch()