add rerun viewer

app.py

@@ -1,14 +1,42 @@
 import gradio as gr
 import spaces
 import torch
+from gradio_rerun import Rerun
+import rerun as rr
+from pathlib import Path
 
-zero = torch.Tensor([0]).cuda()
-print(zero.device) # <-- 'cpu' 🤔
+
+from mini_dust3r.api import OptimizedResult, inferece_dust3r, log_optimized_result
+from mini_dust3r.model import AsymmetricCroCo3DStereo
+
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = AsymmetricCroCo3DStereo.from_pretrained(
+        "naver/DUSt3R_ViTLarge_BaseDecoder_512_dpt"
+    ).to(DEVICE)
 
 @spaces.GPU
-def greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+def predict(image_dir: str):
+    rr.init("my data")
+    optimized_results: OptimizedResult = inferece_dust3r(
+        image_dir=image_dir,
+        model=model,
+        device=DEVICE,
+        batch_size=1,
+    )
+    log_optimized_result(optimized_results, Path("world"))
+    rr.save("dust3r.rrd")
+    return "dust3r.rrd"
+
+with gr.Blocks(css=""".gradio-container {margin: 0 !important; min-width: 100%};""", title="DUSt3R Demo") 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;">DUSt3R Demo</h2>')
+    with gr.Column():
+        inputfiles = gr.File(file_count="multiple")
+        rerun_viewer = Rerun(height=900)
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
+        run_btn = gr.Button("Run")
+        run_btn.click(fn=predict,
+                        inputs=[inputfiles],
+                        outputs=[rerun_viewer])
+        
 demo.launch()