fix ui

app.py

@@ -241,11 +241,10 @@ with gr.Blocks() as demo:
                 presets.keys()) + ["Custom"], label="Choose Preset", value=default_preset)
             nfe_slider = gr.Slider(4, 64, value=len(
                 cache_schedule['attn']), step=1, label="Number of Steps (NFE)")
-            gr.Markdown("", scale=1)
-        with gr.Column(scale=1):
+        with gr.Group():
             gr.Markdown(
-                "Click Grid to Customize Cache Schedule<br>🟧 = Compute Attn Layer / 🟩 = Compute FFN Layer / ⬜ = Cached Layer")
-            image = gr.Image(type="numpy", label=" ", interactive=True, scale=1)
+                "Click Grid to Customize Cache Schedule<br>🟧 = Compute Attn Layer <br> 🟩 = Compute FFN Layer <br> ⬜ = Cached Layer", container=True)
+            image = gr.Image(type="numpy", show_label=False, show_fullscreen_button=False, sources=[], interactive=True, scale=1)
     generate_btn = gr.Button("Synthesize", variant="primary")
     with gr.Row():
         with gr.Group():