Update app.py

app.py

@@ -12,12 +12,12 @@ import spaces
 
 # Model IDs
 MODEL_IDS = {
-    "Model 1 (Widgetcap 448)": "agentsea/paligemma-3b-ft-widgetcap-waveui-448",
-    "Model 2 (WaveUI 896)": "agentsea/paligemma-3b-ft-waveui-896"
+    "paligemma-3b-ft-widgetcap-waveui-448": "agentsea/paligemma-3b-ft-widgetcap-waveui-448",
+    "paligemma-3b-ft-waveui-896": "agentsea/paligemma-3b-ft-waveui-896"
 }
 PROCESSOR_IDS = {
-    "Model 1 (Widgetcap 448)": "google/paligemma-3b-pt-448",
-    "Model 2 (WaveUI 896)": "google/paligemma-3b-pt-896"
+    "paligemma-3b-ft-widgetcap-waveui-448": "google/paligemma-3b-pt-448",
+    "paligemma-3b-ft-waveui-896": "google/paligemma-3b-pt-896"
 }
 COLORS = ['#4285f4', '#db4437', '#f4b400', '#0f9d58', '#e48ef1']
 # Device configuration
@@ -72,11 +72,14 @@ def parse_segmentation(input_image, input_text, model_choice):
 ######## Demo
 
 INTRO_TEXT = """## PaliGemma WaveUI\n\n
-[PaliGemma Widgetcap 448](https://huggingface.co/google/paligemma-3b-ft-widgetcap-448) fine-tuned on the [WaveUI-25k](https://huggingface.co/datasets/agentsea/wave-ui-25k) dataset for UI element detection.\n\n
+
+PaliGemma models fine-tuned on the [WaveUI dataset](https://huggingface.co/datasets/agentsea/wave-ui). There are two base-models:\n\n
+
+- paligemma-3b-ft-widgetcap-waveui-448
+- paligemma-3b-ft-waveui-896
 
 Note:\n\n
 
-- this model is fine-tuned on a subset of the WaveUI dataset and may not generalize to all UI elements.
 - the task it was fine-tuned on was detection, so it may not generalize to other tasks.
 """