Upload app.py

app.py

@@ -166,7 +166,10 @@ def predict(dataset, text, example_index, file, vision_encoder, text_encoder, ts
                 value=f"Please Select Example or Provide CSV File.",
                 visible=True
             ),
-            None
+            None,
+            gr.Markdown(visible=False),  # Hide attention header
+            gr.Gallery(visible=False),   # Hide vision heatmaps
+            gr.Gallery(visible=False)    # Hide time series heatmaps
         )
     elif (vision_encoder is None or text_encoder is None or tsfm is None):
         return (
@@ -174,7 +177,10 @@ def predict(dataset, text, example_index, file, vision_encoder, text_encoder, ts
                 value=f"Please Select Pretrained Model For UniCast.",
                 visible=True
             ),
-            None
+            None,
+            gr.Markdown(visible=False),  # Hide attention header
+            gr.Gallery(visible=False),   # Hide vision heatmaps  
+            gr.Gallery(visible=False)    # Hide time series heatmaps
         )
     else:
         pass
@@ -260,6 +266,7 @@ def predict(dataset, text, example_index, file, vision_encoder, text_encoder, ts
     return (
         gr.Markdown(visible=False),
         fig,
+        gr.Markdown("## Attention Analysis", visible=True), 
         gr.Gallery(vision_heatmap_gallery_items, interactive=False, height="350px", object_fit="contain", visible=True),
         gr.Gallery(time_series_heatmap_gallery_items, interactive=False, height="350px", object_fit="contain", visible=True if time_series_heatmap_gallery_items else False)
     )
@@ -312,13 +319,16 @@ with gr.Blocks() as demo:
     )
     with gr.Row():
         with gr.Column(scale=2):
+            gr.Markdown("## Dataset Selection") 
             dataset_dropdown = gr.Dropdown(["NN5 Daily", "Australian Electricity"], value=None, label="Datasets", interactive=True)
 
             dataset_description_textbox = gr.Textbox(label="Dataset Description", interactive=False)
 
+            gr.Markdown("## Time Series Examples") 
             example_gallery = gr.Gallery(
                 None,
-                interactive=False
+                interactive=False,
+                label="Select Time Series Example" 
             )
             example_index = gr.State(value=None)
             example_gallery.select(selected_example, inputs=example_gallery, outputs=example_index)
@@ -336,14 +346,22 @@ with gr.Blocks() as demo:
             example_index.change(update_time_series_dataframe, inputs=[dataset_dropdown, example_index], outputs=[time_series_file, time_series_dataframe])
 
             time_series_file.change(load_csv, inputs=[example_index, time_series_file], outputs=time_series_dataframe)
+        
         with gr.Column(scale=1):
+            gr.Markdown("## Model Configuration") 
             vision_encoder_radio = gr.Radio(["CLIP", "BLIP"], label="Vision Encoder")
             text_encoder_radio = gr.Radio(["Qwen", "LLaMA"], label="Text Encoder")
             tsfm_radio = gr.Radio(["Timer", "Chronos"], label="Time Series Foundation Model")
             warning_markdown = gr.Markdown(visible=False)
             btn = gr.Button("Run")
+        
         with gr.Column(scale=2):
+            gr.Markdown("## Results")
+
             forecast_plot = gr.Plot(label="Forecast", format="png")
+
+            gr.Markdown("## Attention Analysis", visible=False)
+
             vision_heatmap_gallery = gr.Gallery(visible=False)
             time_series_heatmap_gallery = gr.Gallery(visible=False)