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 = f"{save_dir}/best_image.png"
+        image_path = best_image
         
         if os.path.exists(image_path):
             iter_images = list_iter_images(save_dir)
@@ -132,7 +132,7 @@ with gr.Blocks() as demo:
                 )
             
             with gr.Column():
-                output_image = gr.Image(type="filepath", label="Generated Image")
+                output_image = gr.Image(type="filepath", label="Best Generated Image")
                 status = gr.Textbox(label="Status")
                 iter_gallery = gr.Gallery(label="Iterations", columns=4)