Update app.py

app.py

@@ -1,15 +1,7 @@
 from spaces import GPU
-@GPU
-def dummy_warmup():
-    import torch
-    if torch.cuda.is_available():
-        print("Warmup: GPU is available!")
-        _ = torch.tensor([0.0]).to("cuda")
-dummy_warmup()
-
+import torch
 import numpy as np
 import gradio as gr
-import torch
 import rembg
 import trimesh
 from moge.model.v1 import MoGeModel
@@ -22,6 +14,17 @@ import matplotlib.pyplot as plt
 from eval_wrapper.eval import EvalWrapper, eval_scene
 
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+@GPU
+def dummy_warmup():
+    import torch
+    if torch.cuda.is_available():
+        print("Warmup: GPU is available!")
+        _ = torch.tensor([0.0]).to(device)
+dummy_warmup()
+
+
 outdir = "/tmp/rayst3r"
 
 # loading all necessary models
@@ -101,12 +104,12 @@ def rayst3r_to_glb(img,depth_dict,mask,max_total_points=10e6,rotated=False):
     
     dino_model = torch.hub.load('facebookresearch/dinov2', "dinov2_vitl14_reg")
     dino_model.eval()
-    dino_model.to("cuda")
+    dino_model.to(device)
     
     print("Loading RaySt3R model")
     rayst3r_checkpoint = hf_hub_download("bartduis/rayst3r", "rayst3r.pth")
     rayst3r_model = EvalWrapper(rayst3r_checkpoint,device='cpu')
-    rayst3r_model = rayst3r_model.to("cuda")
+    rayst3r_model = rayst3r_model.to(device)
     
     rayst3r_points = eval_scene(rayst3r_model,os.path.join(outdir, "input"),do_filter_all_masks=True,dino_model=dino_model).cpu()
 
@@ -159,7 +162,7 @@ def input_to_glb(outdir,img,depth_dict,mask,rotated=False):
 @GPU
 def depth_moge(input_img):
     moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl")
-    moge_model.to("cuda")
+    moge_model.to(device)
     input_img_torch = torch.tensor(input_img / 255, dtype=torch.float32, device='cuda').permute(2, 0, 1)
     output = moge_model.infer(input_img_torch).cpu()
     return output