Update app.py

app.py

@@ -14,9 +14,23 @@ def generate_image(prompt):
     
     # 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"{args.save_dir}/{args.task}/{args.model}_{args.prompt}/best_image.png"
+    return f"{save_dir}/best_image.png"
 
 # Create Gradio interface
 iface = gr.Interface(