Merge pull request huggingface#30 from CrazyBoyM/master

feat: support for load weights from local dir

Author: patrickvonplaten

src/controlnet_aux/hed/__init__.py

@@ -106,7 +106,10 @@ def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None
         else:
             filename = filename or "network-bsds500.pth"
 
-        model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+        if os.path.isdir(pretrained_model_or_path):
+            model_path = os.path.join(pretrained_model_or_path, filename)
+        else:
+            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
 
         netNetwork = Network(model_path)
 

src/controlnet_aux/lineart/__init__.py

@@ -104,8 +104,12 @@ def from_pretrained(cls, pretrained_model_or_path, filename=None, coarse_filenam
         filename = filename or "sk_model.pth"
         coarse_filename = coarse_filename or "sk_model2.pth"
 
-        model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
-        coarse_model_path = hf_hub_download(pretrained_model_or_path, coarse_filename, cache_dir=cache_dir)
+        if os.path.isdir(pretrained_model_or_path):
+            model_path = os.path.join(pretrained_model_or_path, filename)
+            coarse_model_path = os.path.join(pretrained_model_or_path, coarse_filename)
+        else:
+            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+            coarse_model_path = hf_hub_download(pretrained_model_or_path, coarse_filename, cache_dir=cache_dir)
 
         model = Generator(3, 1, 3)
         model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

src/controlnet_aux/lineart_anime/__init__.py

@@ -122,7 +122,10 @@ def __init__(self, model):
     def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None):
         filename = filename or "netG.pth"
 
-        model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+        if os.path.isdir(pretrained_model_or_path):
+            model_path = os.path.join(pretrained_model_or_path, filename)
+        else:
+            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
 
         norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)
         net = UnetGenerator(3, 1, 8, 64, norm_layer=norm_layer, use_dropout=False)

src/controlnet_aux/midas/__init__.py

@@ -2,6 +2,7 @@
 import numpy as np
 from PIL import Image
 import torch
+import os 
 
 from huggingface_hub import hf_hub_download
 from einops import rearrange
@@ -22,7 +23,11 @@ def from_pretrained(cls, pretrained_model_or_path, model_type="dpt_hybrid", file
         else:
             filename = filename or "dpt_hybrid-midas-501f0c75.pt"
 
-        model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+        if os.path.isdir(pretrained_model_or_path):
+            model_path = os.path.join(pretrained_model_or_path, filename)
+        else:
+            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+
         return cls(model_type=model_type, model_path=model_path)
 
     def __call__(self, input_image, a=np.pi * 2.0, bg_th=0.1, depth_and_normal=False):

src/controlnet_aux/mlsd/__init__.py

@@ -22,7 +22,10 @@ def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None
         else:
             filename = filename or "mlsd_large_512_fp32.pth"
 
-        model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+        if os.path.isdir(pretrained_model_or_path):
+            model_path = os.path.join(pretrained_model_or_path, filename)
+        else:
+            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
 
         model = MobileV2_MLSD_Large()
         model.load_state_dict(torch.load(model_path), strict=True)

src/controlnet_aux/normalbae/__init__.py

@@ -24,7 +24,11 @@ def __init__(self, model):
     @classmethod
     def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None):
         filename = filename or "scannet.pt"
-        model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+
+        if os.path.isdir(pretrained_model_or_path):
+            model_path = os.path.join(pretrained_model_or_path, filename)
+        else:
+            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
 
         args = types.SimpleNamespace()
         args.mode = 'client'

src/controlnet_aux/open_pose/__init__.py

@@ -60,9 +60,14 @@ def from_pretrained(cls, pretrained_model_or_path, filename=None, hand_filename=
 
             face_pretrained_model_or_path = pretrained_model_or_path
 
-        body_model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
-        hand_model_path = hf_hub_download(pretrained_model_or_path, hand_filename, cache_dir=cache_dir)
-        face_model_path = hf_hub_download(face_pretrained_model_or_path, face_filename, cache_dir=cache_dir)
+        if os.path.isdir(pretrained_model_or_path):
+            body_model_path = os.path.join(pretrained_model_or_path, filename)
+            hand_model_path = os.path.join(pretrained_model_or_path, hand_filename)
+            face_model_path = os.path.join(face_pretrained_model_or_path, face_filename)
+        else:
+            body_model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+            hand_model_path = hf_hub_download(pretrained_model_or_path, hand_filename, cache_dir=cache_dir)
+            face_model_path = hf_hub_download(face_pretrained_model_or_path, face_filename, cache_dir=cache_dir)
 
         body_estimation = Body(body_model_path)
         hand_estimation = Hand(hand_model_path)

src/controlnet_aux/pidi/__init__.py

@@ -17,7 +17,11 @@ def __init__(self, netNetwork):
     @classmethod
     def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None):
         filename = filename or "table5_pidinet.pth"
-        model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+
+        if os.path.isdir(pretrained_model_or_path):
+            model_path = os.path.join(pretrained_model_or_path, filename)
+        else:
+            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
 
         netNetwork = pidinet()
         netNetwork.load_state_dict({k.replace('module.', ''): v for k, v in torch.load(model_path)['state_dict'].items()})

src/controlnet_aux/zoe/__init__.py

@@ -21,7 +21,11 @@ def __init__(self, model):
     def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=None):
         filename = filename or "ZoeD_M12_N.pt"
 
-        model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+        if os.path.isdir(pretrained_model_or_path):
+            model_path = os.path.join(pretrained_model_or_path, filename)
+        else:
+            model_path = hf_hub_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+            
         conf = get_config("zoedepth", "infer")
         model = ZoeDepth.build_from_config(conf)
         model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))['model'])