Changes to app UI.

app.py

@@ -15,13 +15,13 @@ COLUMNS = 5
 def update_default_concepts(prompt):
     default_concepts = {
         "A dog by a tree": ["dog", "grass", "tree", "background"],
-        "A dragon": ["dragon", "sky", "rock", "cloud"],
+        "A man on the beach": ["man", "dirt", "ocean", "sky"],
         "A hot air balloon": ["balloon", "sky", "water", "tree"]
     }
 
     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)
@@ -76,50 +76,42 @@ def process_inputs(prompt, concepts, seed, layer_start_index, timestep_start_ind
 with gr.Blocks(
     css="""
         .container { 
-            max-width: 1400px; 
+            max-width: 1300px; 
             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; }
+        .application {
+            max-width: 1200px;
+        }
         .generated-image {
             display: flex;
             align-items: center;
             justify-content: center;
             height: 100%; /* Ensures full height */
         }
-        .header {
-            display: flex;
-            flex-direction: column;
-        }
+
         .input {
             height: 47px;
         }
         .input-column {
             flex-direction: column;
             gap: 0px;
+            height: 100%;
         }
         .input-column-label {}
-        .gallery {}
+        .gallery {
+            height: 200px;
+        }
         .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;
+        .gallery-container {
+            scrollbar-width: thin;
+            scrollbar-color: grey black;
         }
-
-        /* Show only on screens wider than 768px (adjust as needed) */
+       
+        /* Show only on screens wider than 768px (adjust as needed) 
         @media (min-width: 1024px) {
             .svg-container {
                 min-width: 150px;
@@ -137,124 +129,203 @@ with gr.Blocks(
         }
          @media (min-width: 1530px) {
             .svg-container {
-                min-width: 200px;
+                min-width: 200px; 
                 width: 300px;
                 padding-top: 400px;
             }
         }
 
+        */
+
+        @media (min-width: 1024px) {
+            .svg-container {
+                min-width: 250px;
+            }
+            #concept-attention-callout-svg {
+                width: 250px;
+            }
+        }
+
 
         @media (max-width: 1024px) {
             .svg-container {
+                display: none !important;
+            }
+            #concept-attention-callout-svg {
                 display: none;
             }
         }
 
-    """
-    # ,
-    # elem_classes="container"
-) as demo:
-    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"
-                    )
+        .header {
+            display: flex;
+            flex-direction: column;
+        }
+        #title {
+            font-size: 4.4em;
+            color: #F3B13E;
+            text-align: center;
+            margin: 5px;
+        }
+        #subtitle {
+            font-size: 3.0em;
+            color: #FAE2BA;
+            text-align: center;
+            margin: 5px;
+        }
+        #abstract { 
+            text-align: center; 
+            font-size: 2.0em;
+            color:rgb(219, 219, 219);
+            margin: 5px;
+            margin-top: 10px;
+        }
+        #links {
+            text-align: center;
+            font-size: 2.0em;
+            margin: 5px;
+        }
+        #links a {
+            color: #93B7E9;
+            text-decoration: none;
+        }
 
-                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
-                    )
+        .svg-container {
+            display: flex;
+            justify-content: center;
+            align-items: center;
+        }
 
-                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"
-                    )
+        .caption-label {
+            font-size: 1.15em;
+        }
 
-            with gr.Row(elem_classes="gallery", scale=8):
+        .gallery label {
+            font-size: 1.15em;
+        }
 
-                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"
-                    )
+    """
+) as demo:
 
-                    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.Column(elem_classes="container"):
+
+
+        with gr.Row(elem_classes="container", scale=8):
+            
+            with gr.Column(elem_classes="application-content", scale=10):
+
+                with gr.Row(scale=3, elem_classes="header"):
+                    gr.HTML("""
+                        <h1 id='title'> ConceptAttention </h1>
+                        <h1 id='subtitle'> Visualize Any Concepts in Your Generated Images </h1>
+                        <h1 id='abstract'> Interpret diffusion models with precise, high-quality heatmaps. </h1>
+                        <h1 id='links'> <a href='https://arxiv.org/abs/2502.04320'> Paper </a> | <a href='https://github.com/helblazer811/ConceptAttention'> Code </a> </h1>
+                    """)
+
+                with gr.Row(elem_classes="input-row", scale=2):
+                    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-container", 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",
+                            # scale=4
+                        )
+
+                        cross_attention_gallery = gr.Gallery(
+                            label="Cross Attention", 
+                            show_label=True, 
+                            # columns=3, 
+                            rows=1,
+                            object_fit="contain", 
+                            # height="200px",
+                            elem_classes="gallery",
+                            # scale=4
+                        )
+
+                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=2, min_width=200, elem_classes="svg-column"):
+
+                with gr.Row(scale=8):
+                    gr.HTML("<div></div>")
+
+                with gr.Row(scale=4, elem_classes="svg-container"):
+                    concept_attention_callout_svg = gr.HTML(
+                        "<img src='/gradio_api/file=ConceptAttentionCallout.svg' id='concept-attention-callout-svg'/>",
+                        # container=False,
                     )
 
-            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,
-            )
+                with gr.Row(scale=4):
+                    gr.HTML("<div></div>")
 
 if __name__ == "__main__":
     if os.path.exists("/data-nvme/zerogpu-offload"):