disable loadinug when there is a mask

app.py

@@ -387,14 +387,19 @@ def update_dataset_from_head(head: str) -> Dict[str, Any]:
     return gr.update()
 
 
-def update_upload_component_state(head: str) -> Dict[str, Any]:
+def update_upload_component_state(head: str) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """Disable upload component for heads that require masks."""
     if head in HEADS_REQUIRING_MASK:
-        return gr.update(
-            interactive=False,
-            info="Custom image upload is disabled for this task because it requires a mask from the dataset.",
+        info_update = gr.update(
+            value="⚠️ Custom image upload is disabled for this task because it requires a mask from the dataset.",
+            visible=True,
         )
-    return gr.update(interactive=True, info=None)
+        upload_update = gr.update(interactive=False)
+        return info_update, upload_update
+
+    info_update = gr.update(visible=False)
+    upload_update = gr.update(interactive=True)
+    return info_update, upload_update
 
 
 def load_dataset_metadata(subset: str) -> Tuple[Dict[str, Any], str]:
@@ -557,10 +562,11 @@ def build_demo() -> gr.Blocks:
             """
         )
 
+        default_head = "kits"
         head_dropdown = gr.Dropdown(
             label="Model head",
             choices=[(label, key) for key, label in HEAD_OPTIONS],
-            value="anatomy-ct",
+            value=default_head,
         )
 
         gr.Markdown("---")
@@ -571,7 +577,7 @@ def build_demo() -> gr.Blocks:
                 dataset_dropdown = gr.Dropdown(
                     label="CuriaBench subset",
                     choices=[(meta["label"], key) for key, meta in DATASET_OPTIONS.items()],
-                    value="anatomy-ct",
+                    value="kits",
                 )
                 dataset_status = gr.Markdown("Select a dataset to load class metadata.")
                 class_dropdown = gr.Dropdown(label="Target class filter", choices=["Random"], value="Random")
@@ -579,7 +585,22 @@ def build_demo() -> gr.Blocks:
 
             with gr.Column():
                 gr.Markdown("### Upload custom image")
-                upload_component = gr.Image(label="Upload image", image_mode="L", type="numpy")
+                # Set initial state based on default head
+                initial_requires_mask = default_head in HEADS_REQUIRING_MASK
+                upload_info_text = gr.Markdown(
+                    value=(
+                        "⚠️ Custom image upload is disabled for this task because it requires a mask from the dataset."
+                        if initial_requires_mask
+                        else ""
+                    ),
+                    visible=initial_requires_mask,
+                )
+                upload_component = gr.Image(
+                    label="Upload image",
+                    image_mode="L",
+                    type="numpy",
+                    interactive=not initial_requires_mask,
+                )
 
         gr.Markdown("---")
 
@@ -605,7 +626,7 @@ def build_demo() -> gr.Blocks:
         ).then(
             fn=update_upload_component_state,
             inputs=[head_dropdown],
-            outputs=[upload_component],
+            outputs=[upload_info_text, upload_component],
         )
 
         dataset_dropdown.change(
@@ -628,7 +649,7 @@ def build_demo() -> gr.Blocks:
 
         upload_component.upload(
             fn=handle_upload_preview,
-            inputs=[upload_component],
+            inputs=[upload_component, dataset_dropdown],
             outputs=[
                 image_display,
                 status_text,