clean outputs directories

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from main import main
 from arguments import parse_args
 import os
+import shutil
 import glob
 
 def list_iter_images(save_dir):
@@ -21,6 +22,27 @@ def list_iter_images(save_dir):
 
     return image_paths
 
+def clean_dir(save_dir):
+    # Check if the directory exists
+    if os.path.exists(save_dir):
+        # Check if the directory contains any files
+        if len(os.listdir(save_dir)) > 0:
+            # If it contains files, delete all files in the directory
+            for filename in os.listdir(save_dir):
+                file_path = os.path.join(save_dir, filename)
+                try:
+                    if os.path.isfile(file_path) or os.path.islink(file_path):
+                        os.unlink(file_path)  # Remove file or symbolic link
+                    elif os.path.isdir(file_path):
+                        shutil.rmtree(file_path)  # Remove directory and its contents
+                except Exception as e:
+                    print(f"Failed to delete {file_path}. Reason: {e}")
+            print(f"All files in {save_dir} have been deleted.")
+        else:
+            print(f"{save_dir} exists but is empty.")
+    else:
+        print(f"{save_dir} does not exist.")
+
 def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Progress(track_tqdm=True)):
     # Set up arguments
     args = parse_args()
@@ -32,6 +54,21 @@ def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Pro
     args.cache_dir = "./HF_model_cache"
     args.save_dir = "./outputs"
     args.save_all_images = True
+
+    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}"
+    clean_dir(save_dir)
     
     try:
         # Run the main function with progress tracking
@@ -39,20 +76,6 @@ def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Pro
             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 ''}"
-            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
         image_path = f"{save_dir}/best_image.png"