text2video

Source Code/annotator/midas/__init__.py

@@ -8,13 +8,14 @@ from .api import MiDaSInference
 
 class MidasDetector:
     def __init__(self):
-        self.model = MiDaSInference(model_type="dpt_hybrid").cuda()
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.model = MiDaSInference(model_type="dpt_hybrid").to(self.device)
 
     def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1):
         assert input_image.ndim == 3
         image_depth = input_image
         with torch.no_grad():
-            image_depth = torch.from_numpy(image_depth).float().cuda()
+            image_depth = torch.from_numpy(image_depth).float().to(self.device)
             image_depth = image_depth / 127.5 - 1.0
             image_depth = rearrange(image_depth, 'h w c -> 1 c h w')
             depth = self.model(image_depth)[0]

Source Code/annotator/midas/api.py

@@ -89,8 +89,13 @@ def load_model(model_type):
         normalization = NormalizeImage(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
 
     elif model_type == "dpt_hybrid":  # DPT-Hybrid
+        if not os.path.exists(annotator_ckpts_path):
+            os.makedirs(annotator_ckpts_path)
+            
         if not os.path.exists(model_path):
-            raise FileNotFoundError(f"MiDaS DPT Hybrid model not found at {model_path}. Please ensure dpt_hybrid-midas-501f0c75.pt is present in annotator/ckpts/.")
+            from torch.hub import download_url_to_file
+            print(f"Downloading MiDaS DPT Hybrid from {remote_model_path}...")
+            download_url_to_file(remote_model_path, model_path)
 
         model = DPTDepthModel(
             path=model_path,