update UI

app.py

@@ -204,13 +204,13 @@ MAX_IMAGE_SIZE = 1024
 
 accelerator = accelerate.Accelerator()
 
-def generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps):
+def generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps, use_free_unip):
     """Helper function to generate image with specific number of steps"""
     scheduler = CustomedUniPCMultistepScheduler.from_config(pipe.scheduler.config
                                                             , solver_order = 2 if num_inference_steps==8 else 1
                                                             ,denoise_to_zero = False
                                                             , use_afs = True
-                                                            , use_free_predictor = False)
+                                                            , use_free_predictor = use_free_unip)
     start_free_at_step = 4
     pipe.scheduler = scheduler
     pipe.to('cuda')
@@ -242,7 +242,7 @@ def generate_image_with_steps(prompt, negative_prompt, seed, width, height, guid
                 negative_prompts = 'lowres, bad anatomy, bad hands, watermark'
                 negative_prompts = 1 * [negative_prompts]
                 use_afs = True
-                use_free_predictor = False
+                use_free_predictor = use_free_unip
                 prompt_embeds, cond_kwargs = prepare_sdxl_pipeline_step_parameter(pipe
                                                                                       , prompts
                                                                                       , need_cfg=True
@@ -286,6 +286,7 @@ def infer(
     resolution,
     guidance_scale,
     num_inference_steps,
+    use_free_unip,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -295,7 +296,7 @@ def infer(
     width, height = map(int, resolution.split('x'))
     
     # Generate image with selected steps
-    image_quick = generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps)
+    image_quick = generate_image_with_steps(prompt, negative_prompt, seed, width, height, guidance_scale, num_inference_steps, use_free_unip)
     pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config
                                                              , final_sigmas_type="sigma_min"
                                                              , algorithm_type="sde-dpmsolver++"
@@ -396,21 +397,27 @@ with gr.Blocks() as demo:
                     label="Number of inference steps",
                 )
 
+                use_free_unip = gr.Checkbox(
+                    label="Use free Uni-P predictor",
+                    value=False,
+                )
+
         gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            resolution,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, result_30_steps, seed],
-    )
+        gr.on(
+            triggers=[run_button.click, prompt.submit],
+            fn=infer,
+            inputs=[
+                prompt,
+                negative_prompt,
+                seed,
+                randomize_seed,
+                resolution,
+                guidance_scale,
+                num_inference_steps,
+                use_free_unip,
+            ],
+            outputs=[result, result_30_steps, seed],
+        )
 
 if __name__ == "__main__":
     demo.launch()