Speedup

README.md

@@ -4,7 +4,7 @@ emoji: 🦌
 colorFrom: pink
 colorTo: indigo
 sdk: gradio
-sdk_version: 5.45.0
+sdk_version: 5.1.0
 app_file: app.py
 pinned: false
 short_description: Reward-based Noise Optimization for 1-step t2i models

app.py

@@ -78,36 +78,25 @@ def setup_model(loaded_model_setup, prompt, model, seed, num_iterations, enable_
     args.save_all_images = True
 
     if enable_hps is True:
-        args.enable_hps = True
+        args.disable_hps = False
         args.hps_weighting = hps_w
-    else:
-        args.enable_hps = False
     
     if enable_imagereward is True:
-        args.enable_imagereward = True
+        args.disable_imagereward = False
         args.imagereward_weighting = imgrw_w
-    else:
-        args.enable_imagereward = False
     
     if enable_pickscore is True:
-        args.enable_pickscore = True
+        args.disable_pickscore = False
         args.pickscore_weighting = pcks_w
-    else:
-        args.enable_pickscore = False
     
     if enable_clip is True:
-        args.enable_clip = True
+        args.disable_clip = False
         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__'):
@@ -275,12 +264,7 @@ 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. 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 drastic 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."
+description="Enter a prompt to generate an image using ReNO. Adjust the model and parameters as needed."
 
 css="""
 #model-status-id{
@@ -315,28 +299,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="sdxl-turbo")
+                    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")
-                    learning_rate = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=2.5, label="Learning Rate")
+                    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")
 
                 with gr.Accordion("Advanced Settings", open=True):
                     with gr.Column():
                         with gr.Row():
-                            enable_hps = gr.Checkbox(label="HPS ON", value=True, scale=1)
+                            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=True, scale=1)
+                            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=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)
+                            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=True, scale=1)
+                            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")
@@ -344,11 +328,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",
-                        "A blue scooter is parked near a curb in front of a green vintage car",
+                        "Dwayne Johnson depicted as a philosopher king in an academic painting by Greg Rutkowski",
                         "A curious, orange fox and a fluffy, white rabbit, playing together in a lush, green meadow filled with yellow dandelions",
-                        "An orange chair to the right of a black airplane",
-                        "A brain riding a rocketship towards the moon",
+                        "An epic oil painting: a red portal infront of a cityscape, a solitary figure, and a colorful sky over snowy mountains"
                     ],
                     inputs = [prompt]     
                 )

arguments.py

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

requirements.txt

@@ -1,5 +1,5 @@
-torch==2.5.1
-torchvision==0.20.1
+torch==2.3
+torchvision==0.18.0
 pytorch-lightning==2.2
 datasets==2.18
 transformers==4.38.2
@@ -12,5 +12,4 @@ blobfile
 openai-clip
 setuptools==60.2
 optimum
-xformers
-hf-xet==1.1.8
+xformers