update-params

app.py

@@ -78,25 +78,36 @@ def setup_model(loaded_model_setup, prompt, model, seed, num_iterations, enable_
     args.save_all_images = True
 
     if enable_hps is True:
-        args.disable_hps = False
+        args.enable_hps = True
         args.hps_weighting = hps_w
+    else:
+        args.enable_hps = False
     
     if enable_imagereward is True:
-        args.disable_imagereward = False
+        args.enable_imagereward = True
         args.imagereward_weighting = imgrw_w
+    else:
+        args.enable_imagereward = False
     
     if enable_pickscore is True:
-        args.disable_pickscore = False
+        args.enable_pickscore = True
         args.pickscore_weighting = pcks_w
+    else:
+        args.enable_pickscore = False
     
     if enable_clip is True:
-        args.disable_clip = False
+        args.enable_clip = True
         args.clip_weighting = clip_w
+    else:
+        args.enable_clip = False
 
     if model == "flux":
         args.cpu_offloading = True
         args.enable_multi_apply = True
         args.multi_step_model = "flux"
+    
+    if model == "hyper-sd":
+        args.cpu_offloading = True
 
     # Check if args are the same as the loaded_model_setup except for the prompt
     if loaded_model_setup and hasattr(loaded_model_setup[0], '__dict__'):
@@ -264,7 +275,12 @@ def combined_function(gallery_state, loaded_model_setup, prompt, chosen_model, s
 
 # Create Gradio interface
 title="# ReNO: Enhancing One-step Text-to-Image Models through Reward-based Noise Optimization"
-description="Enter a prompt to generate an image using ReNO. Adjust the model and parameters as needed."
+description = "Enter a prompt to generate an image using ReNO. The method enhances text-to-image generation by optimizing \
+    the initial noise using reward models as detailed in the paper. The demo uses a lower learning rate (2.5) compared to the paper's default (5.0) \
+    for smoother trajectories - if you are looking for more dramatic changes, you can increase this value. You can also \
+    adjust the reward weights to e.g. prioritize either prompt following (increase ImageReward) or aesthetic quality \
+    (increase HPS/PickScore) based on your preferences.\n\nThe first time you load this demo, it will take a bit \
+    to download and initialize the required model. Once loaded, each optimization run takes about 25-60 seconds."
 
 css="""
 #model-status-id{
@@ -299,28 +315,28 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
             with gr.Column():
                 prompt = gr.Textbox(label="Prompt")
                 with gr.Row():
-                    chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd", "flux"], label="Model", value="sd-turbo")
+                    chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd", "flux"], label="Model", value="sdxl-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=10, label="Number of Iterations")
-                    learning_rate = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=5.0, label="Learning Rate")
+                    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=2.5, label="Learning Rate")
 
                 with gr.Accordion("Advanced Settings", open=True):
                     with gr.Column():
                         with gr.Row():
-                            enable_hps = gr.Checkbox(label="HPS ON", value=False, scale=1)
+                            enable_hps = gr.Checkbox(label="HPS ON", value=True, 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)
+                            enable_imagereward = gr.Checkbox(label="ImageReward ON", value=True, 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)
+                            enable_pickscore = gr.Checkbox(label="PickScore ON", value=True, scale=1)
+                            pcks_w = gr.Slider(label="PickScore weight", step=0.01, minimum=0, maximum=0.5, value=0.05, interactive=False, scale=3)
                         with gr.Row():
-                            enable_clip = gr.Checkbox(label="CLIP ON", value=False, scale=1)
+                            enable_clip = gr.Checkbox(label="CLIP ON", value=True, 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")
@@ -328,11 +344,11 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
                 gr.Examples(
                     examples = [
                         "A red dog and a green cat",
-                        "A pink elephant and a grey cow",
-                        "A toaster riding a bike",
-                        "Dwayne Johnson depicted as a philosopher king in an academic painting by Greg Rutkowski",
+                        "A blue scooter is parked near a curb in front of a green vintage car",
                         "A curious, orange fox and a fluffy, white rabbit, playing together in a lush, green meadow filled with yellow dandelions",
-                        "An epic oil painting: a red portal infront of a cityscape, a solitary figure, and a colorful sky over snowy mountains"
+                        "An orange chair to the right of a black airplane"
+                        "A toaster riding a bike",
+                        "A brain riding a rocketship towards the moon",
                     ],
                     inputs = [prompt]     
                 )

arguments.py

@@ -39,16 +39,15 @@ def parse_args():
 
     # reward losses
     parser.add_argument(
-        "--disable_hps", default=True, action="store_false", dest="enable_hps"
+        "--enable_hps", default=False, action="store_true",
     )
     parser.add_argument(
         "--hps_weighting", type=float, help="Weighting for HPS", default=5.0
     )
     parser.add_argument(
-        "--disable_imagereward",
-        default=True,
-        action="store_false",
-        dest="enable_imagereward",
+        "--enable_imagereward",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
         "--imagereward_weighting",
@@ -57,16 +56,15 @@ def parse_args():
         default=1.0,
     )
     parser.add_argument(
-        "--disable_clip", default=True, action="store_false", dest="enable_clip"
+        "--enable_clip", default=False, action="store_true"
     )
     parser.add_argument(
         "--clip_weighting", type=float, help="Weighting for CLIP", default=0.01
     )
     parser.add_argument(
-        "--disable_pickscore",
-        default=True,
-        action="store_false",
-        dest="enable_pickscore",
+        "--enable_pickscore",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
         "--pickscore_weighting",