Add demo model selection to Gradio app

mithridatium/gradio_app.py

@@ -27,6 +27,12 @@ except PackageNotFoundError:
 HIDDEN_DATASET_CHOICES = {"cifar10_for_imagenet", "fake_imagenet"}
 DATASET_CHOICES = sorted(set(utils.DATASET_CONFIGS) - HIDDEN_DATASET_CHOICES)
 UI_DEFENSE_CHOICES = sorted((set(DEFENSES) | {"freeeagle"}) - {"aeva"})
+MODEL_SOURCE_DEMO = "demo"
+MODEL_SOURCE_UPLOAD = "upload"
+DEMO_MODELS = {
+    "clean_resnet18_cifar10": "models/resnet18_cifar10.pt",
+    "backdoored_resnet18_cifar10": "models/resnet18_poison.pt",
+}
 UI_CSS = """
 .gradio-container .run-detection-btn,
 .gradio-container .run-detection-btn button,
@@ -66,6 +72,8 @@ def _write_temp_json(payload: dict[str, Any]) -> str:
 
 
 def _run_detection_from_ui(
+    model_source: str,
+    demo_model: str,
     model_file: str | None,
     dataset: str,
     defense: str,
@@ -78,12 +86,23 @@ def _run_detection_from_ui(
         logs.append(msg)
 
     try:
-        selected_model = (model_file or "").strip()
-        if not selected_model:
-            raise ValueError("Upload a local model checkpoint (.pth or .pt).")
-
-        logs.append("[ui] provider: local upload")
-        logs.append(f"[ui] model source: {selected_model}")
+        source_key = str(model_source or MODEL_SOURCE_DEMO).strip().lower()
+        if source_key == MODEL_SOURCE_DEMO:
+            demo_name = str(demo_model or "").strip()
+            selected_model = DEMO_MODELS.get(demo_name, "")
+            if not selected_model:
+                supported = ", ".join(DEMO_MODELS)
+                raise ValueError(f"Choose a demo model. Supported demo models: {supported}.")
+            if not Path(selected_model).is_file():
+                raise ValueError(f"Demo model is not available in this Space: {selected_model}")
+            logs.append("[ui] model source: demo model")
+            logs.append(f"[ui] demo model: {demo_name}")
+        else:
+            selected_model = (model_file or "").strip()
+            if not selected_model:
+                raise ValueError("Upload a local model checkpoint (.pth or .pt).")
+            logs.append("[ui] model source: uploaded file")
+            logs.append(f"[ui] uploaded model: {selected_model}")
 
         detection = run_detection(
             model=selected_model,
@@ -126,7 +145,7 @@ def build_app() -> gr.Blocks:
     with gr.Blocks(title="Mithridatium UI", css=UI_CSS) as demo:
         gr.Markdown(
             "## Mithridatium Detection UI\n"
-            "Upload a local checkpoint and run a defense."
+            "Choose a demo checkpoint or upload a local checkpoint, then run a defense."
         )
 
         defense = gr.Dropdown(
@@ -135,11 +154,32 @@ def build_app() -> gr.Blocks:
             label="Defense",
         )
 
-        model_file = gr.File(
-            label="Upload Local Model Checkpoint (.pth or .pt)",
-            file_count="single",
-            file_types=[".pth", ".pt"],
-            type="filepath",
+        model_source = gr.State(MODEL_SOURCE_DEMO)
+
+        with gr.Tabs():
+            with gr.Tab("Demo Model") as demo_model_tab:
+                demo_model = gr.Radio(
+                    choices=list(DEMO_MODELS),
+                    value="clean_resnet18_cifar10",
+                    label="Demo Model",
+                )
+            with gr.Tab("Upload Model") as upload_model_tab:
+                model_file = gr.File(
+                    label="Upload Local Model Checkpoint (.pth or .pt)",
+                    file_count="single",
+                    file_types=[".pth", ".pt"],
+                    type="filepath",
+                )
+
+        demo_model_tab.select(
+            fn=lambda: MODEL_SOURCE_DEMO,
+            inputs=None,
+            outputs=model_source,
+        )
+        upload_model_tab.select(
+            fn=lambda: MODEL_SOURCE_UPLOAD,
+            inputs=None,
+            outputs=model_source,
         )
 
         dataset = gr.Dropdown(
@@ -174,6 +214,8 @@ def build_app() -> gr.Blocks:
         run_btn.click(
             fn=_run_detection_from_ui,
             inputs=[
+                model_source,
+                demo_model,
                 model_file,
                 dataset,
                 defense,