Update ppd/models/ppd.py

ppd/models/ppd.py

@@ -6,6 +6,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import cv2
 import random
+from huggingface_hub import hf_hub_download
 from ppd.utils.timesteps import Timesteps
 from ppd.utils.schedule import LinearSchedule
 from ppd.utils.sampler import EulerSampler
@@ -23,7 +24,7 @@ class PixelPerfectDepth(nn.Module):
     ):
         super(PixelPerfectDepth, self).__init__()
 
-        DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
+        DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         self.device = DEVICE
 
         self.semantics_encoder = DepthAnythingV2(
@@ -31,6 +32,10 @@ class PixelPerfectDepth(nn.Module):
             features=256,
             out_channels=[256, 512, 1024, 1024]
         )
+        semantics_pth = hf_hub_download(
+            repo_id="depth-anything/Depth-Anything-V2-Large", 
+            filename="depth_anything_v2_vitl.pth", 
+            repo_type="model")
         self.semantics_encoder.load_state_dict(torch.load(semantics_pth, map_location='cpu'), strict=False)
         self.semantics_encoder = self.semantics_encoder.to(self.device).eval()
         self.dit = DiT()