add weights parameters

app.py

@@ -24,7 +24,7 @@ def list_iter_images(save_dir):
         image_paths.extend(glob.glob(os.path.join(save_dir, f'*.{ext}')))
 
     # Now image_paths contains the list of all image file paths
-    print(image_paths)
+    #print(image_paths)
 
     return image_paths
 
@@ -50,11 +50,12 @@ def clean_dir(save_dir):
         print(f"{save_dir} does not exist.")
 
 def start_over(gallery_state):
+    torch.cuda.empty_cache()  # Free up cached memory
     if gallery_state is not None:
         gallery_state = None
     return gallery_state, None, None, gr.update(visible=False)
 
-def setup_model(prompt, model, num_iterations, learning_rate, progress=gr.Progress(track_tqdm=True)):
+def setup_model(prompt, model, num_iterations, learning_rate, hps_w, imgrw_w, pcks_w, clip_w, progress=gr.Progress(track_tqdm=True)):
 
     """Clear CUDA memory before starting the training."""
     torch.cuda.empty_cache()  # Free up cached memory
@@ -70,14 +71,17 @@ def setup_model(prompt, model, num_iterations, learning_rate, progress=gr.Progre
     args.save_dir = "./outputs"
     args.save_all_images = True
 
+    args.hps_weighting = hps_w
+    args.imagereward_weighting = imgrw_w
+    args.pickscore_weighting = pcks_w
+    args.clip_weighting = clip_w
+
     args, trainer, device, dtype, shape, enable_grad, settings = setup(args)
     loaded_setup = [args, trainer, device, dtype, shape, enable_grad, settings]
 
     return None, loaded_setup
 
 def generate_image(setup_args, num_iterations):
-
-    """Clear CUDA memory before starting executing task."""
     torch.cuda.empty_cache()  # Free up cached memory
 
     args = setup_args[0]
@@ -88,6 +92,7 @@ def generate_image(setup_args, num_iterations):
     enable_grad = setup_args[5]
 
     settings = setup_args[6]
+    print(f"SETTINGS: {settings}")
 
     save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
     clean_dir(save_dir)
@@ -139,7 +144,7 @@ def generate_image(setup_args, num_iterations):
 
     except Exception as e:
         torch.cuda.empty_cache()  # Free up cached memory
-        yield (None, f"An error occurred: {str(e)}", None)
+        return (None, f"An error occurred: {str(e)}", None)
 
 def show_gallery_output(gallery_state):
     if gallery_state is not None:
@@ -179,6 +184,13 @@ with gr.Blocks() as demo:
                     n_iter = gr.Slider(minimum=10, maximum=100, step=10, value=50, label="Number of Iterations")
                     learning_rate = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=5.0, label="Learning Rate")
 
+                with gr.Accordion("Advanced Settings", open=False):
+                    with gr.Column():
+                        hps_w = gr.Slider(label="HPS weight", step=0.1, minimum=0.0, maximum=10.0, value=5.0)
+                        imgrw_w = gr.Slider(label="ImageReward weight", step=0.1, minimum=0, maximum=5.0, value=1.0)
+                        pcks_w = gr.Slider(label="PickScore weight", step=0.01, minimum=0, maximum=5.0, value=0.05)
+                        clip_w = gr.Slider(label="CLIP weight", step=0.01, minimum=0, maximum=0.1, value=0.01)
+
                 submit_btn = gr.Button("Submit")
 
                 gr.Examples(
@@ -204,7 +216,7 @@ with gr.Blocks() as demo:
         outputs = [gallery_state, output_image, status, iter_gallery]
     ).then(
         fn = setup_model,
-        inputs = [prompt, chosen_model, n_iter, learning_rate],
+        inputs = [prompt, chosen_model, n_iter, hps_w, imgrw_w, pcks_w, clip_w, learning_rate],
         outputs =  [output_image, loaded_model_setup]
     ).then(
         fn = generate_image,