Update app.py

app.py

@@ -22,6 +22,7 @@ def get_temp_path():
 
 @spaces.GPU(duration=120)
 def scene_generate(rgb,num_coarse_views,num_mcs_views,mcs_rect_w,mcs_steps):
+    torch.cuda.init()
     # coarse
     vistadream.scene = Gaussian_Scene(cfg)
     # for trajectory genearation
@@ -39,6 +40,9 @@ def scene_generate(rgb,num_coarse_views,num_mcs_views,mcs_rect_w,mcs_steps):
     # coarse scene
     vistadream._coarse_scene(rgb)
     torch.cuda.empty_cache()
+
+@spaces.GPU(duration=120)
+def scene_refinement()
     # refinement
     vistadream._MCS_Refinement()
     output_path = get_temp_path()
@@ -52,7 +56,8 @@ def render_video(output_path):
     return output_path+'video_rgb.mp4',output_path+'video_dpt.mp4'
 
 def process(rgb,num_coarse_views,num_mcs_views,mcs_rect_w,mcs_steps):
-    path = scene_generate(rgb,num_coarse_views,num_mcs_views,mcs_rect_w,mcs_steps)
+    scene_generate(rgb,num_coarse_views,num_mcs_views,mcs_rect_w,mcs_steps)
+    path = scene_refinement()
     rgb.save(output_path+'input.png')
     return render_video(path)