Update app.py

app.py

@@ -559,7 +559,7 @@ def draw_mask(o, b, v, d, evt: gr.EventData):
 
 
 load_model="""
-async(c, o, p, d, n, m)=>{
+async(c, o, p, d, n, m, s)=>{
   var intv = setInterval(function(){
     if (document.getElementById("model3D").getElementsByTagName("canvas")[0]) {
       try {
@@ -664,10 +664,14 @@ async(c, o, p, d, n, m)=>{
           BABYLON.Engine.LastCreatedScene.meshes[BABYLON.Engine.LastCreatedScene.meshes.length-1].rotation.y = coords[o_[j]].heading / 180 * Math.PI;
           BABYLON.Engine.LastCreatedScene.meshes[BABYLON.Engine.LastCreatedScene.meshes.length-1].rotation.z = -coords[o_[j]].pitch / 180 * Math.PI;
         }*/
-          
+
+        if (s == false) {
+          v_url = document.getElementById("output_video").getElementsByTagName("video")[0].src;
+        } else {
+          v_url = document.getElementById("depth_video").getElementsByTagName("video")[0].src;
+        }
         window.videoDome = new BABYLON.VideoDome(
-            "videoDome",
-            [document.getElementById("output_video").getElementsByTagName("video")[0].src],
+            "videoDome", [v_url],
             {
                 resolution: 16,
                 size: 512,
@@ -1367,7 +1371,7 @@ with gr.Blocks(css=css, js=js, head=head) as demo:
         return output_video_path + (json.dumps(locations),)
 
     submit.click(on_submit, inputs=[input_video, model_type, blur_in, boffset, bsize, coords], outputs=[processed_video, processed_zip, output_frame, output_mask, output_depth, depth_video, coords])
-    render.click(None, inputs=[coords, mesh_order, output_frame, output_mask, selected, output_depth], outputs=None, js=load_model)
+    render.click(None, inputs=[coords, mesh_order, output_frame, output_mask, selected, output_depth, output_switch], outputs=None, js=load_model)
     render.click(partial(get_mesh), inputs=[output_frame, output_mask, blur_in, load_all], outputs=[result, mesh_order])
 
     example_files = [["./examples/streetview.mp4", "vits", blurin, 1, 32, example_coords]]