Update app.py

app.py

@@ -78,7 +78,7 @@ def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Pro
         best_image, total_init_rewards, total_best_rewards = main(args, progress_callback)
         
         # Return the path to the generated image
-        image_path = best_image
+        image_path = f"{save_dir}/best_image.png"
         
         if os.path.exists(image_path):
             iter_images = list_iter_images(save_dir)