Update app.py

app.py

@@ -14,8 +14,8 @@ def load_description(fp):
 
 
 @spaces.GPU(duration=90)
-def run_theia(image, pred_iou_thresh, stability_score_thresh):
-    theia_model = AutoModel.from_pretrained("theaiinstitute/theia-tiny-patch16-224-cddsv", trust_remote_code=True)
+def run_theia(model_size, image, pred_iou_thresh, stability_score_thresh):
+    theia_model = AutoModel.from_pretrained(f"theaiinstitute/theia-{model_size}-patch16-224-cddsv", trust_remote_code=True)
     theia_model = theia_model.to('cuda')
     target_model_names = [
         "google/vit-huge-patch14-224-in21k",
@@ -72,6 +72,7 @@ with gr.Blocks() as demo:
             input_image = gr.Image(label="Input Image", type="pil")
 
             with gr.Accordion("Advanced Settings", open=False):
+                model_size = gr.Radio(["tiny", "small", "base"], value="tiny", label="Theia Model Size")
                 pred_iou_thresh = gr.Slider(0.05, 0.95, step=0.05, value=0.5, label="SAM Pred IoU Thresh")
                 stability_score_thresh = gr.Slider(0.05, 0.95, step=0.05, value=0.7, label="SAM Stability Score Thresh")
 
@@ -84,7 +85,7 @@ with gr.Blocks() as demo:
 
     submit_button.click(
         run_theia,
-        inputs=[input_image, pred_iou_thresh, stability_score_thresh],
+        inputs=[model_size, input_image, pred_iou_thresh, stability_score_thresh],
         outputs=[dinov2_output, sam_output, depth_anything_output]
     )