bp

app.py

@@ -17,17 +17,10 @@ from spaces import GPU
 outdir = "/tmp/rayst3r"
 
 # loading all necessary models
-print("Loading DINOv2 model")
-dino_model = torch.hub.load('facebookresearch/dinov2', "dinov2_vitl14_reg")
-dino_model.eval()
 
 print("Loading MoGe model")
 # Load the model from huggingface hub (or load from local).
-moge_model = MoGeModel.from_pretrained("Ruicheng/moge-vitl")
 
-print("Loading RaySt3R model")
-rayst3r_checkpoint = hf_hub_download("bartduis/rayst3r", "rayst3r.pth")
-rayst3r_model = EvalWrapper(rayst3r_checkpoint,device='cpu')
 
 def depth2uint16(depth):
     return depth * torch.iinfo(torch.uint16).max / 10.0 # threshold is in m, convert to uint16 value
@@ -98,8 +91,14 @@ def prep_for_rayst3r(img,depth_dict,mask):
 def rayst3r_to_glb(img,depth_dict,mask,max_total_points=10e6,rotated=False):
     prep_for_rayst3r(img,depth_dict,mask)
     
+    dino_model = torch.hub.load('facebookresearch/dinov2', "dinov2_vitl14_reg")
+    dino_model.eval()
     dino_model.to("cuda")
-    rayst3r_model.to("cuda")
+    
+    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_points = eval_scene(rayst3r_model,os.path.join(outdir, "input"),do_filter_all_masks=True,dino_model=dino_model).cpu()
 
@@ -151,6 +150,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")
     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()

eval_wrapper/eval.py

@@ -38,6 +38,9 @@ class EvalWrapper(torch.nn.Module):
         self.model.load_state_dict(checkpoint['model'])
         self.model.eval()
 
+    def to(self,device):
+        self.model.to(device)
+
     def forward(self,x,dino_model=None):
         pred, gt, loss, scale = eval_model(self.model,x,mode='viz',dino_model=dino_model,return_scale=True)
         return pred, gt, loss, scale