UI changes, add option to activate weighting types

app.py

@@ -57,7 +57,7 @@ def start_over(gallery_state, loaded_model_setup):
         loaded_model_setup = None  # Reset loaded model setup to prevent re-triggering old state
     return gallery_state, None, None, gr.update(visible=False), loaded_model_setup
 
-def setup_model(prompt, model, seed, num_iterations, learning_rate, hps_w, imgrw_w, pcks_w, clip_w, progress=gr.Progress(track_tqdm=True)):
+def setup_model(prompt, model, seed, num_iterations, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate, progress=gr.Progress(track_tqdm=True)):
     if prompt is None or prompt == "":
         raise gr.Error("You forgot to provide a prompt !")
         
@@ -76,10 +76,21 @@ def setup_model(prompt, model, seed, num_iterations, learning_rate, hps_w, imgrw
     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
+    if enable_hps is True:
+        args.disable_hps = False
+        args.hps_weighting = hps_w
+    
+    if enable_imagereward is True:
+        args.disable_imagereward = False
+        args.imagereward_weighting = imgrw_w
+    
+    if enable_pickscore is True:
+        args.disable_pickscore = False
+        args.pickscore_weighting = pcks_w
+    
+    if enable_clip is True:
+        args.disable_clip = False
+        args.clip_weighting = clip_w
 
     if model == "flux":
         args.cpu_offloading = True
@@ -237,6 +248,8 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
                 with gr.Row():
                     chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd", "flux"], label="Model", value="sd-turbo")
                     seed = gr.Number(label="seed", value=0)
+
+                model_status = gr.Textbox(label="model status", visible=True, elem_id="model-status-id")
                 
                 with gr.Row():
                     n_iter = gr.Slider(minimum=10, maximum=100, step=10, value=50, label="Number of Iterations")
@@ -244,10 +257,18 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
 
                 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)
+                        with gr.Row():
+                            enable_hps = gr.Checkbox(label="HPS ON", value=False, scale=1)
+                            hps_w = gr.Slider(label="HPS weight", step=0.1, minimum=0.0, maximum=10.0, value=5.0, interactive=False, scale=3)
+                        with gr.Row():
+                            enable_imagereward = gr.Checkbox(label="ImageReward ON", value=False, scale=1)
+                            imgrw_w = gr.Slider(label="ImageReward weight", step=0.1, minimum=0, maximum=5.0, value=1.0, interactive=False, scale=3)
+                        with gr.Row():
+                            enable_pickscore = gr.Checkbox(label="PickScore ON", value=False, scale=1)
+                            pcks_w = gr.Slider(label="PickScore weight", step=0.01, minimum=0, maximum=5.0, value=0.05, interactive=False, scale=3)
+                        with gr.Row():
+                            enable_clip = gr.Checkbox(label="CLIP ON", value=False, scale=1)
+                            clip_w = gr.Slider(label="CLIP weight", step=0.01, minimum=0, maximum=0.1, value=0.01, interactive=False, scale=3)
 
                 submit_btn = gr.Button("Submit")
 
@@ -264,18 +285,49 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
                 )
             
             with gr.Column():
-                model_status = gr.Textbox(label="model status", visible=True, elem_id="model-status-id")
                 output_image = gr.Image(type="filepath", label="Best Generated Image")
                 status = gr.Textbox(label="Status")
                 iter_gallery = gr.Gallery(label="Iterations", columns=4, visible=False)
 
+    def allow_weighting(weight_type):
+        if weight_type is True:
+            return gr.update(interactive=True)
+        else:
+            return gr.update(interactive=False)
+    
+    enable_hps.change(
+        fn = allow_weighting,
+        inputs = [enable_hps],
+        outputs = [hps_w],
+        queue = False
+    )
+    enable_imagereward.change(
+        fn = allow_weighting,
+        inputs = [enable_imagereward],
+        outputs = [imgrw_w],
+        queue = False
+    )
+    enable_pickscore.change(
+        fn = allow_weighting,
+        inputs = [enable_pickscore],
+        outputs = [pcks_w],
+        queue = False
+    )
+    enable_clip.change(
+        fn = allow_weighting,
+        inputs = [enable_clip],
+        outputs = [clip_w],
+        queue = False
+    )
+    
+    
     submit_btn.click(
         fn = start_over,
         inputs =[gallery_state, loaded_model_setup],  # Reset loaded model setup as well
         outputs = [gallery_state, output_image, status, iter_gallery, loaded_model_setup]  # Ensure loaded_model_setup is reset
     ).then(
         fn = setup_model,
-        inputs = [prompt, chosen_model, seed, n_iter, hps_w, imgrw_w, pcks_w, clip_w, learning_rate],
+        inputs = [prompt, chosen_model, seed, n_iter, enable_hps, hps_w, enable_imagereward, imgrw_w, enable_pickscore, pcks_w, enable_clip, clip_w, learning_rate],
         outputs = [model_status, loaded_model_setup]  # Load the new setup into the state
     ).then(
         fn = generate_image,