try display current image live

app.py

@@ -74,6 +74,8 @@ def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Pro
         # Run the main function with progress tracking
         def progress_callback(step):
             progress(step / num_iterations, f"Iteration {step}/{num_iterations}")
+            current_image = f"{save_dir}/{step}.png"
+            yield current_image, f"Iteration {step}/{num_iterations}", None
 
         best_image, total_init_rewards, total_best_rewards = main(args, progress_callback)