chore: Refactor predict function to handle both single and multiple image inputs

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+
 import spaces
 import torch
 from gradio_rerun import Rerun
@@ -17,11 +18,20 @@ model = AsymmetricCroCo3DStereo.from_pretrained(
 
 
 @spaces.GPU
-def predict(image_name_list: list[str]):
+def predict(image_name_list: list[str] | str):
+    # check if is list or string and if not raise error
+    if not isinstance(image_name_list, list) and not isinstance(image_name_list, str):
+        raise gr.Error(
+            f"Input must be a list of strings or a string, got: {type(image_name_list)}"
+        )
     uuid_str = str(uuid.uuid4())
     filename = Path(f"/tmp/gradio/{uuid_str}.rrd")
     rr.init(f"{uuid_str}")
     log_path = Path("world")
+
+    if isinstance(image_name_list, str):
+        image_name_list = [image_name_list]
+
     optimized_results: OptimizedResult = inferece_dust3r(
         image_dir_or_list=image_name_list,
         model=model,
@@ -41,11 +51,22 @@ with gr.Blocks(
 ) as demo:
     # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
     gr.HTML('<h2 style="text-align: center;">Mini-DUSt3R Demo</h2>')
-    with gr.Column():
-        inputfiles = gr.File(file_count="multiple")
-        rerun_viewer = Rerun(height=900)
+    with gr.Tab(label="Single Image"):
+        with gr.Column():
+            singe_image = gr.Image(type="filepath")
+            run_btn_single = gr.Button("Run")
+            rerun_viewer_single = Rerun(height=900)
+            run_btn_single.click(
+                fn=predict, inputs=[singe_image], outputs=[rerun_viewer_single]
+            )
+    with gr.Tab(label="Multi Image"):
+        with gr.Column():
+            multi_files = gr.File(file_count="multiple")
+            run_btn_multi = gr.Button("Run")
+            rerun_viewer_multi = Rerun(height=900)
+            run_btn_multi.click(
+                fn=predict, inputs=[multi_files], outputs=[rerun_viewer_multi]
+            )
 
-        run_btn = gr.Button("Run")
-        run_btn.click(fn=predict, inputs=[inputfiles], outputs=[rerun_viewer])
 
 demo.launch()