chore: update ui

app.py

@@ -40,6 +40,7 @@ def infer(
     guidance_scale: float,
     num_inference_steps: int,
     tkg_channels: list[int] = [0, 1, 1, 0],
+    chroma_key_shift: float = 0.11,
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -60,7 +61,7 @@ def infer(
     )
     tkg_latents = apply_tkg_noise(
         latents,
-        shift=0.11,
+        shift=chroma_key_shift,
         delta_shift=0.1,
         std_dev=0.5,
         factor=8,
@@ -69,8 +70,8 @@ def infer(
 
     latents = torch.cat(
         [
-            latents,
             tkg_latents,
+            latents,
         ],
         dim=0,
     )
@@ -107,6 +108,7 @@ def on_generate(
     guidance_scale: float,
     num_inference_steps: int,
     color_name: str,
+    chroma_key_shift: float,
     *args,
     **kwargs
 ):
@@ -123,6 +125,7 @@ def on_generate(
         guidance_scale,
         num_inference_steps,
         tkg_channels=tkg_channels,
+        chroma_key_shift=chroma_key_shift,
         *args,
         **kwargs,
     )
@@ -131,7 +134,7 @@ def on_generate(
 
 examples = [
     # "1girl, arima kana, oshi no ko, hoshimachi suisei, hoshimachi suisei \(1st costume\), cosplay, looking at viewer, smile, outdoors, night, v, masterpiece, high score, great score, absurdres",
-    "1girl, solo, upper body, looking at viewer, straight-on, masterpiece, best quality",
+    "1girl, solo, school uniform, cat ears, full body, looking at viewer, straight-on, chibi, masterpiece, best quality",
 ]
 
 
@@ -149,26 +152,36 @@ TKG-DM🥚🍚: Training-free Chroma Key Content Generation Diffusion Model
 
         with gr.Row():
             with gr.Column():
-                prompt = gr.Text(
+                prompt = gr.Textbox(
                     label="Prompt",
                     max_lines=4,
                     placeholder="Enter your prompt",
                 )
-                negative_prompt = gr.Textbox(
-                    label="Negative prompt",
-                    max_lines=4,
-                    placeholder="Enter a negative prompt",
-                    value="lowres, bad anatomy, bad hands, text, error, missing finger, extra digits, fewer digits, cropped, worst quality, low quality, low score, bad score, average score, signature, watermark, username, blurry",
-                )
-
+                
                 color_set = gr.Dropdown(
-                    label="Chroma key color",
+                    label="Background color",
                     choices=list(COLOR_SET_MAP.keys()),
                     value="green",
                 )
 
+                with gr.Accordion("TKG Settings", open=False):
+                    chroma_key_shift = gr.Slider(
+                        label="Latent mean shift for chroma key",
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.05,
+                        value=0.11,  
+                    )
+
 
                 with gr.Accordion("Advanced Settings", open=False):
+                    negative_prompt = gr.Textbox(
+                        label="Negative prompt",
+                        max_lines=4,
+                        placeholder="Enter a negative prompt",
+                        value="lowres, bad anatomy, bad hands, text, error, missing finger, extra digits, fewer digits, cropped, worst quality, low quality, low score, bad score, average score, signature, watermark, username, blurry",
+                    )
+
                     seed = gr.Slider(
                         label="Seed",
                         minimum=0,
@@ -215,8 +228,9 @@ TKG-DM🥚🍚: Training-free Chroma Key Content Generation Diffusion Model
 
             with gr.Column():
                 run_button = gr.Button("Generate", variant="primary")
-                result_w_tkg = gr.Image(label="with TKG")
-                result_wo_tkg = gr.Image(label="without TKG")
+                with gr.Row():
+                    result_w_tkg = gr.Image(label="with TKG")
+                    result_wo_tkg = gr.Image(label="without TKG")
 
         
 
@@ -235,6 +249,7 @@ TKG-DM🥚🍚: Training-free Chroma Key Content Generation Diffusion Model
             guidance_scale,
             num_inference_steps,
             color_set,
+            chroma_key_shift,
         ],
         outputs=[result_w_tkg, result_wo_tkg, seed],
     )