Speedup

README.md

@@ -3,9 +3,8 @@ title: ReNO
 emoji: 🦌
 colorFrom: pink
 colorTo: indigo
-python_version: '3.10'
 sdk: gradio
-sdk_version: 6.14.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]     
                 )
@@ -368,29 +352,25 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
         fn = allow_weighting,
         inputs = [enable_hps],
         outputs = [hps_w],
-        queue = False,
-        api_visibility="private"
+        queue = False
     )
     enable_imagereward.change(
         fn = allow_weighting,
         inputs = [enable_imagereward],
         outputs = [imgrw_w],
-        queue = False,
-        api_visibility="private"
+        queue = False
     )
     enable_pickscore.change(
         fn = allow_weighting,
         inputs = [enable_pickscore],
         outputs = [pcks_w],
-        queue = False,
-        api_visibility="private"
+        queue = False
     )
     enable_clip.change(
         fn = allow_weighting,
         inputs = [enable_clip],
         outputs = [clip_w],
-        queue = False,
-        api_visibility="private"
+        queue = False
     )
 
     submit_btn.click(
@@ -402,8 +382,7 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
         ],
         outputs = [
             gallery_state, output_image, status, iter_gallery, loaded_model_setup, model_status  # Ensure `model_status` is included in the outputs
-        ],
-        api_visibility="private"
+        ]
     )
     
     """
@@ -427,4 +406,4 @@ with gr.Blocks(css=css, analytics_enabled=False) as demo:
     """
 
 # Launch the app
-demo.queue().launch(show_error=True)
+demo.queue().launch(show_error=True, show_api=False)

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,22 +1,15 @@
-torch==2.5.1
-torchvision==0.20.1
-pytorch-lightning==2.2.0
-datasets==2.18.0
-
-transformers==4.55.4
-diffusers==0.35.1
-accelerate==1.8.1
-huggingface_hub==0.34.4
-safetensors>=0.4.3
-
+torch==2.3
+torchvision==0.18.0
+pytorch-lightning==2.2
+datasets==2.18
+transformers==4.38.2
+diffusers
 hpsv2==1.2
 hpsv2x==1.2.0
 image-reward==1.5
-open-clip-torch==2.24.0
+open-clip-torch==2.24
 blobfile
 openai-clip
+setuptools==60.2
 optimum
-xformers
-hf-xet==1.1.8
-
-setuptools>=68
+xformers