Update app.py

app.py

@@ -52,9 +52,9 @@ ckpt_path = hf_hub_download(
 )
 state_dict = torch.load(ckpt_path, map_location="cpu")
 model.load_state_dict(state_dict, strict=False)
-model = model.to(DEVICE).eval()
+model = model.eval()
 
-moge_model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal").to(DEVICE).eval()
+moge_model = MoGeModel.from_pretrained("Ruicheng/moge-2-vitl-normal").eval()
 
 
 def main(share=True):
@@ -68,6 +68,7 @@ def main(share=True):
 
     @spaces.GPU
     def predict_depth(image, denoise_steps):
+        model = model.to(DEVICE)
         depth, resize_image = model.infer_image(image, sampling_steps=denoise_steps)
         return depth, resize_image
 
@@ -75,6 +76,7 @@ def main(share=True):
     def predict_moge_depth(image):
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         image = torch.tensor(image / 255, dtype=torch.float32, device=DEVICE).permute(2, 0, 1)
+        moge_model = moge_model.to(DEVICE)
         metric_depth, mask, intrinsics = moge_model.infer(image)
         metric_depth[~mask] = metric_depth[mask].max()
         return metric_depth, mask, intrinsics