change layout

app.py

@@ -35,6 +35,7 @@ class GradioDemo:
             num_inference_steps = 4,
             sd_pipe_guidance_scale = 1.0,
             seed = 1024,
+            oms_prompt_flag=True,
         ):
         pipe_kwargs = dict(
             prompt = prompt,
@@ -51,18 +52,29 @@ class GradioDemo:
         )['images'][0]
         
         generator = torch.Generator(device=self.pipe.device).manual_seed(seed)
-        pipe_kwargs.update(oms_flag=True, oms_prompt=oms_prompt, oms_guidance_scale=1.0)
-        print(f'w/ oms wo/ cfg kwargs: {pipe_kwargs}')
-        image_oms = self.pipe(
+        pipe_kwargs.update(oms_flag=True, oms_prompt=prompt, oms_guidance_scale=1.0)
+        print(f'w/ oms wo/ cfg (consistent) kwargs: {pipe_kwargs}')
+        image_oms_cp = self.pipe(
             **pipe_kwargs,
             generator=generator
         )['images'][0]
         
+        if oms_prompt_flag:
+            generator = torch.Generator(device=self.pipe.device).manual_seed(seed)
+            pipe_kwargs.update(oms_prompt=oms_prompt)
+            print(f'w/ oms wo/ cfg (inconsistent) kwargs: {pipe_kwargs}')
+            image_oms_icp = self.pipe(
+                **pipe_kwargs,
+                generator=generator
+            )['images'][0]
+        else:
+            image_oms_icp = None
+            
         oms_guidance_flag = oms_guidance_scale != 1.0
         if oms_guidance_flag:
             generator = torch.Generator(device=self.pipe.device).manual_seed(seed)
             pipe_kwargs.update(oms_guidance_scale=oms_guidance_scale)
-            print(f'w/ oms +cfg kwargs: {pipe_kwargs}')
+            print(f'w/ oms +cfg (inconsistent) kwargs: {pipe_kwargs}')
             image_oms_cfg = self.pipe(
                 **pipe_kwargs,
                 generator=generator
@@ -70,29 +82,44 @@ class GradioDemo:
         else:
             image_oms_cfg = None
         
-        return image_raw, image_oms, image_oms_cfg, gr.update(visible=oms_guidance_flag)
+        return image_raw, image_oms_cp, image_oms_icp, image_oms_cfg, gr.update(visible=oms_prompt_flag), gr.update(visible=oms_guidance_flag)
     
     def mainloop(self):
         with gr.Blocks() as demo:
-            gr.Markdown("# One More Step Demo")
+            gr.Markdown("# One More Step for SDXL w/ LCM-LoRA")
+            
+            with gr.Group() as inputs:
+                prompt = gr.Textbox(label="Prompt", value="a cat against orange ground, studio")
+                with gr.Accordion('OMS Prompt'):
+                    oms_prompt_checkbox = gr.Checkbox(info="Inconsistent OMS prompt allows the additional control of low freq info, default is the same as Prompt.", label="Adding OMS Prompt", value=True)
+                    oms_prompt = gr.Textbox(label="OMS Prompt", value="a black cat", info='try "a black cat" and "a black room" for diverse control.')
+                with gr.Accordion('OMS Guidance'):
+                    oms_cfg_scale_checkbox = gr.Checkbox(info="OMS Guidance will enhance the OMS prompt, specially focus on color and brightness. ", label="Adding OMS Guidance", value=True)
+                    oms_guidance_scale = gr.Slider(label="OMS Guidance Scale", minimum=1.0, maximum=5.0, value=2., step=0.1)
+                run_button = gr.Button(value="Generate images")
+                with gr.Accordion("Advanced options", open=False):
+                    num_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=4, step=1)
+                    sd_guidance_scale = gr.Slider(label="SD Pipe Guidance Scale", minimum=1, maximum=3, value=1.0, step=0.1)
+                    seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=False, value=1024)
             with gr.Row():
-                with gr.Column():
-                    prompt = gr.Textbox(label="Prompt", value="a cat against black ground, studio")
-                    oms_prompt = gr.Textbox(label="OMS Prompt", value="an orange cat")
-                    oms_guidance_scale = gr.Slider(label="OMS Guidance Scale", minimum=1.0, maximum=5.0, value=1.5, step=0.1)
-                    run_button = gr.Button(value="Generate images")
-                    with gr.Accordion("Advanced options", open=False):
-                        num_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=4, step=1)
-                        sd_guidance_scale = gr.Slider(label="SD Pipe Guidance Scale", minimum=0.1, maximum=30.0, value=1.0, step=0.1)
-                        seed = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, randomize=False, value=1024)
-                with gr.Column():
-                    output_raw = gr.Image(label="SDXL w/ LCM-LoRA w/o OMS ")
-                    output_oms = gr.Image(label="w/ OMS w/o OMS CFG")
-                    with gr.Column(visible=False) as oms_cfg_wd:
-                        output_oms_cfg = gr.Image(label="w/ OMS w/ OMS CFG")
+                output_raw = gr.Image(label="SDXL w/ LCM-LoRA ")
+                output_oms_cp = gr.Image(label="w/ OMS (consistent prompt) w/o OMS CFG")
+                output_oms_icp = gr.Image(label="w/ OMS (inconsistent prompt) w/o OMS CFG")
+                output_oms_cfg = gr.Image(label="w/ OMS w/ OMS CFG")
 
-            ips = [prompt, oms_prompt, oms_guidance_scale, num_steps, sd_guidance_scale, seed]
-            run_button.click(fn=self._inference, inputs=ips, outputs=[output_raw, output_oms, output_oms_cfg, oms_cfg_wd])
+            oms_prompt_checkbox.input(
+                fn=lambda oms_prompt_flag, prompt, oms_prompt: (oms_prompt if oms_prompt_flag else prompt, gr.update(interactive=oms_prompt_flag)), 
+                inputs=[oms_prompt_checkbox, prompt, oms_prompt], 
+                outputs=[oms_prompt, oms_prompt]
+            )
+            oms_cfg_scale_checkbox.input(
+                fn=lambda oms_cfg_scale_flag: (1.5 if oms_cfg_scale_flag else 1.0, gr.update(interactive=oms_cfg_scale_flag)), 
+                inputs=[oms_cfg_scale_checkbox], 
+                outputs=[oms_guidance_scale, oms_guidance_scale]
+            )
+            
+            ips = [prompt, oms_prompt, oms_guidance_scale, num_steps, sd_guidance_scale, seed, oms_prompt_checkbox]
+            run_button.click(fn=self._inference, inputs=ips, outputs=[output_raw, output_oms_cp, output_oms_icp, output_oms_cfg, output_oms_icp, output_oms_cfg])
 
             demo.queue(max_size=20)
             demo.launch()