fix: dtype, feat: better ui

app.py

@@ -56,6 +56,7 @@ def infer(
         ),
         generator=generator,
         device=device,
+        dtype=torch.bfloat16,
     )
     latents = apply_tkg_noise(
         latents,
@@ -103,7 +104,7 @@ def on_generate(
     tkg_channels = color_name_to_channels(color_name)
     # TODO: custom channels
 
-    return infer(
+    w_tkg, wo_tkg, seed = infer(
         prompt,
         negative_prompt,
         seed,
@@ -113,8 +114,12 @@ def on_generate(
         guidance_scale,
         num_inference_steps,
         tkg_channels=tkg_channels,
+        *args,
+        **kwargs,
     )
 
+    return w_tkg, wo_tkg, seed
+
 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",
@@ -134,71 +139,79 @@ TKG-DM🥚🍚: Training-free Chroma Key Content Generation Diffusion Model
 """)
 
         with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
-
-            run_button = gr.Button("Run", scale=0, variant="primary")
-
-        with gr.Row():
-            result_w_tkg = gr.Image(label="Result", show_label=False)
-            result_wo_tkg = gr.Image(label="Result", show_label=False)
-
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                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,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=832,  
+            with gr.Column():
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter your prompt",
+                    container=False,
                 )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1152,
+                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",
                 )
 
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=5.0,  
+                color_set = gr.Dropdown(
+                    label="Chroma key color",
+                    choices=list(COLOR_SET_MAP.keys()),
+                    value="green",
                 )
 
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=25,  
-                )
+
+                with gr.Accordion("Advanced Settings", open=False):
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0,
+                    )
+
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+                    with gr.Row():
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=832,  
+                        )
+
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=32,
+                            value=1152,
+                        )
+
+                    with gr.Row():
+                        guidance_scale = gr.Slider(
+                            label="Guidance scale",
+                            minimum=0.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=5.0,  
+                        )
+
+                        num_inference_steps = gr.Slider(
+                            label="Number of inference steps",
+                            minimum=1,
+                            maximum=50,
+                            step=1,
+                            value=25,  
+                        )
+
+            with gr.Column():
+                run_button = gr.Button("Run", scale=0, variant="primary")
+                result_w_tkg = gr.Image(label="Result", show_label=False)
+                result_wo_tkg = gr.Image(label="Result", show_label=False)
+
+        
 
         gr.Examples(examples=examples, inputs=[prompt])
 
@@ -214,6 +227,7 @@ TKG-DM🥚🍚: Training-free Chroma Key Content Generation Diffusion Model
             height,
             guidance_scale,
             num_inference_steps,
+            color_set,
         ],
         outputs=[result_w_tkg, result_wo_tkg, seed],
     )