Update app.py

app.py

@@ -16,7 +16,7 @@ from eval_wrapper.eval import EvalWrapper, eval_scene
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-@GPU
+@GPU(duration = 180)
 def dummy_warmup():
     import torch
     if torch.cuda.is_available():
@@ -98,7 +98,7 @@ def prep_for_rayst3r(img,depth_dict,mask):
     # save image
     save_tensor_as_png(torch.from_numpy(img), os.path.join(input_dir, "rgb.png"))
 
-@GPU
+@GPU(duration = 180)
 def rayst3r_to_glb(img,depth_dict,mask,max_total_points=10e6,rotated=False):
     prep_for_rayst3r(img,depth_dict,mask)
     
@@ -159,15 +159,15 @@ def input_to_glb(outdir,img,depth_dict,mask,rotated=False):
     scene.export(outfile)
     return outfile
 
-@GPU
+@GPU(duration = 180)
 def depth_moge(input_img):
     moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl")
     moge_model.to(device)
-    input_img_torch = torch.tensor(input_img / 255, dtype=torch.float32, device='cuda').permute(2, 0, 1)
+    input_img_torch = torch.tensor(input_img / 255, dtype=torch.float32, device=device).permute(2, 0, 1)
     output = moge_model.infer(input_img_torch).cpu()
     return output 
 
-@GPU
+@GPU(duration = 180)
 def mask_rembg(input_img):
     #masked_img = rembg.remove(input_img,)
     output_img = rembg.remove(input_img, alpha_matting=False, post_process_mask=True)
@@ -186,7 +186,7 @@ def mask_rembg(input_img):
     rgb = output_np[:,:,:3]
     return mask, rgb 
 
-@GPU
+@GPU(duration = 180)
 def process_image(input_img):
     # resize the input image
     rotated = False