gradio can track iterations callbacks for single task

app.py

@@ -3,7 +3,7 @@ from main import main
 from arguments import parse_args
 import os
 
-def generate_image(prompt, model, num_iterations, learning_rate, progress = gr.Progress(track_tqdm=True)):
+def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Progress(track_tqdm=True)):
     # Set up arguments
     args = parse_args()
     args.task = "single"
@@ -16,8 +16,11 @@ def generate_image(prompt, model, num_iterations, learning_rate, progress = gr.P
     args.save_all_images = True
     
     try:
-        # Run the main function
-        main(args)
+        # Run the main function with progress tracking
+        def progress_callback(step):
+            progress(step / num_iterations, f"Iteration {step}/{num_iterations}")
+
+        best_image, total_init_rewards, total_best_rewards = main(args, progress_callback)
 
         settings = (
             f"{args.model}{'_' + args.prompt if args.task == 't2i-compbench' else ''}"
@@ -66,7 +69,7 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 prompt = gr.Textbox(label="Prompt")
-                chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd"], label="Model")
+                chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd"], label="Model", value="sd-turbo")
                 
                 with gr.Row():
                     n_iter = gr.Slider(minimum=10, maximum=100, step=10, value=50, label="Number of Iterations")