try without partial

app.py

@@ -10,16 +10,13 @@ import sys
 import os.path as path
 import torch
 import tempfile
+import gradio
 
 HERE_PATH = path.normpath(path.dirname(__file__))  # noqa
 MASt3R_REPO_PATH = path.normpath(path.join(HERE_PATH, './mast3r'))  # noqa
 sys.path.insert(0, MASt3R_REPO_PATH)  # noqa
 
-import mast3r.demo
-mast3r.demo.get_reconstructed_scene = spaces.GPU(mast3r.demo.get_reconstructed_scene)
-mast3r.demo.get_3D_model_from_scene = spaces.GPU(mast3r.demo.get_3D_model_from_scene)
-
-from mast3r.demo import main_demo
+from mast3r.demo import get_reconstructed_scene, get_3D_model_from_scene, set_scenegraph_options
 from mast3r.model import AsymmetricMASt3R
 from mast3r.utils.misc import hash_md5
 
@@ -35,9 +32,138 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model = AsymmetricMASt3R.from_pretrained(weights_path).to(device)
 chkpt_tag = hash_md5(weights_path)
 
-# mast3r will write the 3D model inside tmpdirname/chkpt_tag
-with tempfile.TemporaryDirectory(suffix='_mast3r_gradio_demo') as tmpdirname:
-    cache_path = os.path.join(tmpdirname, chkpt_tag)
-    os.makedirs(cache_path, exist_ok=True)
-    main_demo(tmpdirname, model, device, 512, server_name=None, server_port=None,
-              silent=True, share=None, gradio_delete_cache=7200)
+tmpdirname = "tmp/gradio"
+image_size = 512
+silent = True
+gradio_delete_cache = 7200
+
+
+@spaces.GPU()
+def local_get_reconstructed_scene(current_scene_state,
+                                  filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,
+                                  as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize,
+                                  win_cyclic, refid, TSDF_thresh, shared_intrinsics, **kw):
+    return get_reconstructed_scene(tmpdirname, gradio_delete_cache, model, device, silent, image_size, current_scene_state,
+                                   filelist, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,
+                                   as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size, scenegraph_type, winsize,
+                                   win_cyclic, refid, TSDF_thresh, shared_intrinsics, **kw)
+
+
+@spaces.GPU()
+def local_get_3D_model_from_scene(scene_state, min_conf_thr=2, as_pointcloud=False, mask_sky=False,
+                                  clean_depth=False, transparent_cams=False, cam_size=0.05, TSDF_thresh=0):
+    return get_3D_model_from_scene(silent, scene_state, min_conf_thr, as_pointcloud, mask_sky,
+                                   clean_depth, transparent_cams, cam_size, TSDF_thresh)
+
+
+recon_fun = local_get_reconstructed_scene
+model_from_scene_fun = local_get_3D_model_from_scene
+
+
+def get_context(delete_cache):
+    css = """.gradio-container {margin: 0 !important; min-width: 100%};"""
+    title = "MASt3R Demo"
+    if delete_cache:
+        return gradio.Blocks(css=css, title=title, delete_cache=(delete_cache, delete_cache))
+    else:
+        return gradio.Blocks(css=css, title="MASt3R Demo")  # for compatibility with older versions
+
+
+with get_context(gradio_delete_cache) as demo:
+    # scene state is save so that you can change conf_thr, cam_size... without rerunning the inference
+    scene = gradio.State(None)
+    gradio.HTML('<h2 style="text-align: center;">MASt3R Demo</h2>')
+    with gradio.Column():
+        inputfiles = gradio.File(file_count="multiple")
+        with gradio.Row():
+            with gradio.Column():
+                with gradio.Row():
+                    lr1 = gradio.Slider(label="Coarse LR", value=0.07, minimum=0.01, maximum=0.2, step=0.01)
+                    niter1 = gradio.Number(value=500, precision=0, minimum=0, maximum=10_000,
+                                           label="num_iterations", info="For coarse alignment!")
+                    lr2 = gradio.Slider(label="Fine LR", value=0.014, minimum=0.005, maximum=0.05, step=0.001)
+                    niter2 = gradio.Number(value=200, precision=0, minimum=0, maximum=100_000,
+                                           label="num_iterations", info="For refinement!")
+                    optim_level = gradio.Dropdown(["coarse", "refine", "refine+depth"],
+                                                  value='refine', label="OptLevel",
+                                                  info="Optimization level")
+                with gradio.Row():
+                    matching_conf_thr = gradio.Slider(label="Matching Confidence Thr", value=5.,
+                                                      minimum=0., maximum=30., step=0.1,
+                                                      info="Before Fallback to Regr3D!")
+                    shared_intrinsics = gradio.Checkbox(value=False, label="Shared intrinsics",
+                                                        info="Only optimize one set of intrinsics for all views")
+                    scenegraph_type = gradio.Dropdown([("complete: all possible image pairs", "complete"),
+                                                       ("swin: sliding window", "swin"),
+                                                       ("logwin: sliding window with long range", "logwin"),
+                                                       ("oneref: match one image with all", "oneref")],
+                                                      value='complete', label="Scenegraph",
+                                                      info="Define how to make pairs",
+                                                      interactive=True)
+                    with gradio.Column(visible=False) as win_col:
+                        winsize = gradio.Slider(label="Scene Graph: Window Size", value=1,
+                                                minimum=1, maximum=1, step=1)
+                        win_cyclic = gradio.Checkbox(value=False, label="Cyclic sequence")
+                    refid = gradio.Slider(label="Scene Graph: Id", value=0,
+                                          minimum=0, maximum=0, step=1, visible=False)
+        run_btn = gradio.Button("Run")
+
+        with gradio.Row():
+            # adjust the confidence threshold
+            min_conf_thr = gradio.Slider(label="min_conf_thr", value=1.5, minimum=0.0, maximum=10, step=0.1)
+            # adjust the camera size in the output pointcloud
+            cam_size = gradio.Slider(label="cam_size", value=0.2, minimum=0.001, maximum=1.0, step=0.001)
+            TSDF_thresh = gradio.Slider(label="TSDF Threshold", value=0., minimum=0., maximum=1., step=0.01)
+        with gradio.Row():
+            as_pointcloud = gradio.Checkbox(value=True, label="As pointcloud")
+            # two post process implemented
+            mask_sky = gradio.Checkbox(value=False, label="Mask sky")
+            clean_depth = gradio.Checkbox(value=True, label="Clean-up depthmaps")
+            transparent_cams = gradio.Checkbox(value=False, label="Transparent cameras")
+
+        outmodel = gradio.Model3D()
+
+        # events
+        scenegraph_type.change(set_scenegraph_options,
+                               inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
+                               outputs=[win_col, winsize, win_cyclic, refid])
+        inputfiles.change(set_scenegraph_options,
+                          inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
+                          outputs=[win_col, winsize, win_cyclic, refid])
+        win_cyclic.change(set_scenegraph_options,
+                          inputs=[inputfiles, win_cyclic, refid, scenegraph_type],
+                          outputs=[win_col, winsize, win_cyclic, refid])
+        run_btn.click(fn=recon_fun,
+                      inputs=[scene, inputfiles, optim_level, lr1, niter1, lr2, niter2, min_conf_thr, matching_conf_thr,
+                              as_pointcloud, mask_sky, clean_depth, transparent_cams, cam_size,
+                              scenegraph_type, winsize, win_cyclic, refid, TSDF_thresh, shared_intrinsics],
+                      outputs=[scene, outmodel])
+        min_conf_thr.release(fn=model_from_scene_fun,
+                             inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                     clean_depth, transparent_cams, cam_size, TSDF_thresh],
+                             outputs=outmodel)
+        cam_size.change(fn=model_from_scene_fun,
+                        inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                clean_depth, transparent_cams, cam_size, TSDF_thresh],
+                        outputs=outmodel)
+        TSDF_thresh.change(fn=model_from_scene_fun,
+                           inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                   clean_depth, transparent_cams, cam_size, TSDF_thresh],
+                           outputs=outmodel)
+        as_pointcloud.change(fn=model_from_scene_fun,
+                             inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                     clean_depth, transparent_cams, cam_size, TSDF_thresh],
+                             outputs=outmodel)
+        mask_sky.change(fn=model_from_scene_fun,
+                        inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                clean_depth, transparent_cams, cam_size, TSDF_thresh],
+                        outputs=outmodel)
+        clean_depth.change(fn=model_from_scene_fun,
+                           inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                   clean_depth, transparent_cams, cam_size, TSDF_thresh],
+                           outputs=outmodel)
+        transparent_cams.change(model_from_scene_fun,
+                                inputs=[scene, min_conf_thr, as_pointcloud, mask_sky,
+                                        clean_depth, transparent_cams, cam_size, TSDF_thresh],
+                                outputs=outmodel)
+demo.launch(share=None, server_name=None, server_port=None)