download from huggingfacehub

README.md

@@ -14,6 +14,7 @@ from controlnet_aux import OpenposeDetector, MLSDdetector, HEDdetector, CannyDet
 open_pose = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
 mlsd = MLSDdetector.from_pretrained("lllyasviel/ControlNet")
 hed = HEDdetector.from_pretrained("lllyasviel/ControlNet")
+midas = MidasDetector.from_pretrained("lllyasviel/ControlNet")
 canny = CannyDetector()
-midas = MidasDetector()
+
 ```

src/controlnet_aux/midas/__init__.py

@@ -3,14 +3,22 @@
 from PIL import Image
 import torch
 
+from huggingface_hub import hf_hub_download
 from einops import rearrange
 from .api import MiDaSInference
 
 
 class MidasDetector:
-    def __init__(self):
-        self.model = MiDaSInference(model_type="dpt_hybrid").cuda()
+    def __init__(self, model_type="dpt_hybrid", model_path=None):
+        self.model = MiDaSInference(model_type=model_type, model_path=model_path).cuda()
 
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path, model_type="dpt_hybrid", filename=None):
+        filename = filename or "annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
+        model_path = hf_hub_download(pretrained_model_or_path, filename)
+        return cls(model_type=model_type, model_path=model_path)
+
     def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1):
 
         input_type = "np"

src/controlnet_aux/midas/api.py

@@ -74,10 +74,10 @@ def load_midas_transform(model_type):
     return transform
 
 
-def load_model(model_type):
+def load_model(model_type, model_path=None):
     # https://github.com/isl-org/MiDaS/blob/master/run.py
     # load network
-    model_path = ISL_PATHS[model_type]
+    model_path = model_path or ISL_PATHS[model_type]
     if model_type == "dpt_large":  # DPT-Large
         model = DPTDepthModel(
             path=model_path,
@@ -155,10 +155,10 @@ class MiDaSInference(nn.Module):
         "midas_v21_small",
     ]
 
-    def __init__(self, model_type):
+    def __init__(self, model_type, model_path):
         super().__init__()
         assert (model_type in self.MODEL_TYPES_ISL)
-        model, _ = load_model(model_type)
+        model, _ = load_model(model_type, model_path)
         self.model = model
         self.model.train = disabled_train