Changes to UI

app.py

@@ -12,24 +12,6 @@ from concept_attention import ConceptAttentionFluxPipeline
 IMG_SIZE = 210
 COLUMNS = 5
 
-EXAMPLES = [
-    [
-        "A dog by a tree",  # prompt
-        "tree, dog, grass, background",  # words
-        42,  # seed
-    ],
-    # [
-    #     "A dragon",  # prompt
-    #     "dragon, sky, rock, cloud",  # words
-    #     42,  # seed
-    # ],
-    # [
-    #     "A hot air balloon",  # prompt
-    #     "balloon, sky, water, tree",  # words
-    #     42,  # seed
-    # ]
-]
-
 def update_default_concepts(prompt):
     default_concepts = {
         "A dog by a tree": ["dog", "grass", "tree", "background"],
@@ -39,7 +21,7 @@ def update_default_concepts(prompt):
 
     return gr.update(value=default_concepts.get(prompt, []))
 
-pipeline = ConceptAttentionFluxPipeline(model_name="flux-schnell", device="cuda") #, offload_model=True)
+pipeline = ConceptAttentionFluxPipeline(model_name="flux-schnell", device="cuda:2", offload_model=True)
 
 def convert_pil_to_bytes(img):
     img = img.resize((IMG_SIZE, IMG_SIZE), resample=Image.NEAREST)
@@ -51,21 +33,14 @@ def convert_pil_to_bytes(img):
 
 @spaces.GPU(duration=60)
 def process_inputs(prompt, concepts, seed, layer_start_index, timestep_start_index):
-    # print("Processing inputs")
-    # assert layer_start_index is not None
-    # assert timestep_start_index is not None
+    if not prompt:
+        raise gr.exceptions.InputError("prompt", "Please enter a prompt")
 
     if not prompt.strip():
         raise gr.exceptions.InputError("prompt", "Please enter a prompt")
 
     prompt = prompt.strip()
 
-    print(concepts)
-    # if not word_list.strip():
-    #     gr.exceptions.InputError("words", "Please enter comma-separated words")
-
-    # concepts = [w.strip() for w in word_list.split(",")]
-
     if len(concepts) == 0:
         raise gr.exceptions.InputError("words", "Please enter at least 1 concept")
     
@@ -94,18 +69,17 @@ def process_inputs(prompt, concepts, seed, layer_start_index, timestep_start_ind
     cross_attention_heatmaps = [heatmap.resize((IMG_SIZE, IMG_SIZE), resample=Image.NEAREST) for heatmap in cross_attention_heatmaps]
     cross_attention_maps_and_labels = [(cross_attention_heatmaps[concept_index], concepts[concept_index]) for concept_index in range(len(concepts))]
 
-    # heatmaps_and_labels = [(concept_heatmaps[concept_index], concepts[concept_index]) for concept_index in range(len(concepts))]
-    # all_images_and_labels = [(output_image, "Generated Image")] + heatmaps_and_labels
-    # num_rows = math.ceil(len(all_images_and_labels) / COLUMNS)
-
     return output_image, \
         gr.update(value=output_space_maps_and_labels, columns=len(output_space_maps_and_labels)), \
         gr.update(value=cross_attention_maps_and_labels, columns=len(cross_attention_maps_and_labels))
 
 with gr.Blocks(
     css="""
-        .container { max-width: 1200px; margin: 0 auto; padding: 20px; }
-        .title { text-align: center; margin-bottom: 10px; }
+        .container { 
+            max-width: 1400px; 
+            margin: 0 auto; 
+            padding: 20px; 
+        }
         .authors { text-align: center; margin-bottom: 10px; }
         .affiliations { text-align: center; color: #666; margin-bottom: 10px; }
         .abstract { text-align: center; margin-bottom: 40px; }
@@ -115,6 +89,10 @@ with gr.Blocks(
             justify-content: center;
             height: 100%; /* Ensures full height */
         }
+        .header {
+            display: flex;
+            flex-direction: column;
+        }
         .input {
             height: 47px;
         }
@@ -123,159 +101,167 @@ with gr.Blocks(
             gap: 0px;
         }
         .input-column-label {}
-        .gallery {
-            # scrollbar-width: thin;
-            # scrollbar-color: #27272A;
-        }
-
+        .gallery {}
         .run-button-column {
             width: 100px !important;
         }
+        #title {
+            font-size: 2.4em;
+            text-align: center;
+            margin-bottom: 10px;
+        }
+        #subtitle {
+            font-size: 2.0em;
+            text-align: center;
+        }
+
+        #concept-attention-callout-svg {
+            width: 250px;
+        }
+
+        /* Show only on screens wider than 768px (adjust as needed) */
+        @media (min-width: 1024px) {
+            .svg-container {
+                min-width: 150px;
+                width: 200px;
+                padding-top: 540px;
+            }
+        }
+
+        @media (min-width: 1280px) {
+            .svg-container {
+                min-width: 200px;
+                width: 300px;
+                padding-top: 400px;
+            }
+        }
+         @media (min-width: 1530px) {
+            .svg-container {
+                min-width: 200px;
+                width: 300px;
+                padding-top: 370px;
+            }
+        }
+
+
+        @media (max-width: 1024px) {
+            .svg-container {
+                display: none;
+            }
+        }
+
     """
     # ,
     # elem_classes="container"
 ) as demo:
-    with gr.Column(elem_classes="container"):
-        gr.Markdown("# ConceptAttention: Visualize Any Concepts in Your Generated Images", elem_classes="title")
-        # gr.Markdown("### Alec Helbling¹, Tuna Meral², Ben Hoover¹³, Pinar Yanardag², Duen Horng (Polo) Chau¹", elem_classes="authors")
-        # gr.Markdown("### ¹Georgia Tech · ²Virginia Tech · ³IBM Research", elem_classes="affiliations")
-        gr.Markdown("## Interpret generative models with precise, high-quality heatmaps. Check out our paper [here](https://arxiv.org/abs/2502.04320).", elem_classes="abstract")
-
-        with gr.Row(scale=1, equal_height=True):
-            with gr.Column(scale=3, elem_classes="input-column"):
-                gr.HTML(
-                    "Write a Prompt",
-                    elem_classes="input-column-label"
-                )
-                prompt = gr.Dropdown(
-                    ["A dog by a tree", "A dragon", "A hot air balloon"], 
-                    # label="Prompt", 
-                    container=False,
-                    # scale=3,
-                    allow_custom_value=True,
-                    elem_classes="input"
-                )
-
-            with gr.Column(scale=7, elem_classes="input-column"):
-                gr.HTML(
-                    "Select or Write Concepts",
-                    elem_classes="input-column-label"
-                )
-                concepts = gr.Dropdown(
-                    ["dog", "grass", "tree", "dragon", "sky", "rock", "cloud", "balloon", "water", "background"], 
-                    value=["dog", "grass", "tree", "background"], 
-                    multiselect=True, 
-                    label="Concepts",
-                    container=False,
-                    allow_custom_value=True,
-                    # scale=4,
-                    elem_classes="input",
-                    max_choices=5
-                )
-
-            with gr.Column(scale=1, min_width=100, elem_classes="input-column run-button-column"):
-                gr.HTML(
-                    "​",
-                    elem_classes="input-column-label"
-                )
-                submit_btn = gr.Button(
-                    "Run",
-                    # scale=1,
-                    elem_classes="input"
-                )
-            # prompt = gr.Textbox(
-            #     label="Enter your prompt", 
-            #     placeholder="Enter your prompt", 
-            #     value=EXAMPLES[0][0],
-            #     scale=4,
-            #     # show_label=True,
-            #     container=False
-            #     # height="80px"
-            # )
-            # words = gr.Textbox(
-            #     label="Enter a list of concepts (comma-separated)", 
-            #     placeholder="Enter a list of concepts (comma-separated)", 
-            #     value=EXAMPLES[0][1],
-            #     scale=4,
-            #     # show_label=True,
-            #     container=False
-            #     # height="80px"
-            # )
-
-        num_rows_state = gr.State(value=1)  # Initial number of rows
-
-        # generated_image = gr.Image(label="Generated Image", elem_classes="input-image")
-        # gallery = gr.Gallery(
-        #     label="Generated images", 
-        #     show_label=True, 
-        #     # elem_id="gallery",
-        #     columns=COLUMNS, 
-        #     rows=1,
-        #     # object_fit="contain", 
-        #     height="auto",
-        #     elem_classes="gallery"
-        # )
-
-        with gr.Row(elem_classes="gallery", scale=8):
-
-            with gr.Column(scale=1):
-                generated_image = gr.Image(
-                    elem_classes="generated-image",
-                    show_label=False
-                )
-                
-            with gr.Column(scale=4):
-                concept_attention_gallery = gr.Gallery(
-                    label="Concept Attention (Ours)", 
-                    show_label=True, 
-                    # columns=3, 
-                    rows=1,
-                    object_fit="contain", 
-                    height="200px",
-                    elem_classes="gallery"
-                )
-
-                cross_attention_gallery = gr.Gallery(
-                    label="Cross Attention", 
-                    show_label=True, 
-                    # columns=3, 
-                    rows=1,
-                    object_fit="contain", 
-                    height="200px",
-                    elem_classes="gallery"
-                )
-
-        with gr.Accordion("Advanced Settings", open=False):
-            seed = gr.Slider(minimum=0, maximum=10000, step=1, label="Seed", value=42)
-            layer_start_index = gr.Slider(minimum=0, maximum=18, step=1, label="Layer Start Index", value=10)
-            timestep_start_index = gr.Slider(minimum=0, maximum=4, step=1, label="Timestep Start Index", value=2)
-
-        submit_btn.click(
-            fn=process_inputs, 
-            inputs=[prompt, concepts, seed, layer_start_index, timestep_start_index], 
-            outputs=[generated_image, concept_attention_gallery, cross_attention_gallery]
-        )
-
-        # gr.Examples(examples=EXAMPLES, inputs=[prompt, concepts, seed, layer_start_index, timestep_start_index], outputs=[gallery, num_rows_state], fn=process_inputs, cache_examples=False)
-        # num_rows_state.change(
-        #     fn=lambda rows: gr.Gallery.update(rows=int(rows)),
-        #     inputs=[num_rows_state],
-        #     outputs=[gallery]
-        # )
-
-        prompt.change(update_default_concepts, inputs=[prompt], outputs=[concepts])
-
-        # Automatically process the first example on launch
-        demo.load(
-            process_inputs, 
-            inputs=[prompt, concepts, seed, layer_start_index, timestep_start_index], 
-            outputs=[generated_image, concept_attention_gallery, cross_attention_gallery]
-        )
+    with gr.Row(elem_classes="container"):
+        with gr.Column(elem_classes="application", scale=15):
+            with gr.Row(scale=3, elem_classes="header"):
+                gr.HTML("<h1 id='title'> ConceptAttention: Visualize Any Concepts in Your Generated Images</h1>")
+                gr.HTML("<h2 id='subtitle'> Interpret generative models with precise, high-quality heatmaps. <br/> Check out our paper <a href='https://arxiv.org/abs/2502.04320'> here </a>. </h2>")
+
+            with gr.Row(scale=1, equal_height=True):
+                with gr.Column(scale=4, elem_classes="input-column", min_width=250):
+                    gr.HTML(
+                        "Write a Prompt",
+                        elem_classes="input-column-label"
+                    )
+                    prompt = gr.Dropdown(
+                        ["A dog by a tree", "A dragon", "A hot air balloon"], 
+                        container=False,
+                        allow_custom_value=True,
+                        elem_classes="input"
+                    )
+
+                with gr.Column(scale=7, elem_classes="input-column"):
+                    gr.HTML(
+                        "Select or Write Concepts",
+                        elem_classes="input-column-label"
+                    )
+                    concepts = gr.Dropdown(
+                        ["dog", "grass", "tree", "dragon", "sky", "rock", "cloud", "balloon", "water", "background"], 
+                        value=["dog", "grass", "tree", "background"], 
+                        multiselect=True, 
+                        label="Concepts",
+                        container=False,
+                        allow_custom_value=True,
+                        # scale=4,
+                        elem_classes="input",
+                        max_choices=5
+                    )
+
+                with gr.Column(scale=1, min_width=100, elem_classes="input-column run-button-column"):
+                    gr.HTML(
+                        "&#8203;",
+                        elem_classes="input-column-label"
+                    )
+                    submit_btn = gr.Button(
+                        "Run",
+                        elem_classes="input"
+                    )
+
+            with gr.Row(elem_classes="gallery", scale=8):
+
+                with gr.Column(scale=1, min_width=250):
+                    generated_image = gr.Image(
+                        elem_classes="generated-image",
+                        show_label=False
+                    )
+                    
+                with gr.Column(scale=4):
+                    concept_attention_gallery = gr.Gallery(
+                        label="Concept Attention (Ours)", 
+                        show_label=True, 
+                        # columns=3, 
+                        rows=1,
+                        object_fit="contain", 
+                        height="200px",
+                        elem_classes="gallery",
+                        elem_id="concept-attention-gallery"
+                    )
+
+                    cross_attention_gallery = gr.Gallery(
+                        label="Cross Attention", 
+                        show_label=True, 
+                        # columns=3, 
+                        rows=1,
+                        object_fit="contain", 
+                        height="200px",
+                        elem_classes="gallery"
+                    )
+
+            with gr.Accordion("Advanced Settings", open=False):
+                seed = gr.Slider(minimum=0, maximum=10000, step=1, label="Seed", value=42)
+                layer_start_index = gr.Slider(minimum=0, maximum=18, step=1, label="Layer Start Index", value=10)
+                timestep_start_index = gr.Slider(minimum=0, maximum=4, step=1, label="Timestep Start Index", value=2)
+
+            submit_btn.click(
+                fn=process_inputs, 
+                inputs=[prompt, concepts, seed, layer_start_index, timestep_start_index], 
+                outputs=[generated_image, concept_attention_gallery, cross_attention_gallery]
+            )
+
+            prompt.change(update_default_concepts, inputs=[prompt], outputs=[concepts])
+
+            # Automatically process the first example on launch
+            demo.load(
+                process_inputs, 
+                inputs=[prompt, concepts, seed, layer_start_index, timestep_start_index], 
+                outputs=[generated_image, concept_attention_gallery, cross_attention_gallery]
+            )
+
+        with gr.Column(scale=4, min_width=250, elem_classes="svg-container"):
+            concept_attention_callout_svg = gr.HTML(
+                "<img src='/gradio_api/file=ConceptAttentionCallout.svg' id='concept-attention-callout-svg'/>",
+                # container=False,
+            )
 
 if __name__ == "__main__":
     if os.path.exists("/data-nvme/zerogpu-offload"):
         subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
-    demo.launch()
+    demo.launch(
+        allowed_paths=["."]
+    )
     #     share=True,
     #     server_name="0.0.0.0",
     #     inbrowser=True,