app.py

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.cache_dir = "./HF_model_cache"
    args.save_dir = "./outputs"
    args.save_all_images = True
    
    # Run the main function
    main(args)

    settings = (
        f"{args.model}{'_' + args.prompt if args.task == 't2i-compbench' else ''}"
        f"{'_no-optim' if args.no_optim else ''}_{args.seed if args.task != 'geneval' else ''}"
        f"_lr{args.lr}_gc{args.grad_clip}_iter{args.n_iters}"
        f"_reg{args.reg_weight if args.enable_reg else '0'}"
        f"{'_pickscore' + str(args.pickscore_weighting) if args.enable_pickscore else ''}"
        f"{'_clip' + str(args.clip_weighting) if args.enable_clip else ''}"
        f"{'_hps' + str(args.hps_weighting) if args.enable_hps else ''}"
        f"{'_imagereward' + str(args.imagereward_weighting) if args.enable_imagereward else ''}"
        f"{'_aesthetic' + str(args.aesthetic_weighting) if args.enable_aesthetic else ''}"
    )

    save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
    
    # Return the path to the generated image
    return f"{save_dir}/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()