CoolFace
Apppublic

fffiloni/ControlNet-Video

sourceHugging Faceupdated 3y agoView on Hugging Face
428likes
patch116 linesDownload Raw Back to root
1diff --git a/annotator/hed/__init__.py b/annotator/hed/__init__.py2index 42d8dc6..1587035 1006443--- a/annotator/hed/__init__.py4+++ b/annotator/hed/__init__.py5@@ -1,8 +1,12 @@6+import pathlib7+8 import numpy as np9 import cv210 import torch11 from einops import rearrange12 13+root_dir = pathlib.Path(__file__).parents[2]14+15 16 class Network(torch.nn.Module):17     def __init__(self):18@@ -64,7 +68,7 @@ class Network(torch.nn.Module):19             torch.nn.Sigmoid()20         )21 22-        self.load_state_dict({strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.load('./annotator/ckpts/network-bsds500.pth').items()})23+        self.load_state_dict({strKey.replace('module', 'net'): tenWeight for strKey, tenWeight in torch.load(f'{root_dir}/annotator/ckpts/network-bsds500.pth').items()})24     # end25 26     def forward(self, tenInput):27diff --git a/annotator/midas/api.py b/annotator/midas/api.py28index 9fa305e..d8594ea 10064429--- a/annotator/midas/api.py30+++ b/annotator/midas/api.py31@@ -1,5 +1,7 @@32 # based on https://github.com/isl-org/MiDaS33 34+import pathlib35+36 import cv237 import torch38 import torch.nn as nn39@@ -10,10 +12,11 @@ from .midas.midas_net import MidasNet40 from .midas.midas_net_custom import MidasNet_small41 from .midas.transforms import Resize, NormalizeImage, PrepareForNet42 43+root_dir = pathlib.Path(__file__).parents[2]44 45 ISL_PATHS = {46-    "dpt_large": "annotator/ckpts/dpt_large-midas-2f21e586.pt",47-    "dpt_hybrid": "annotator/ckpts/dpt_hybrid-midas-501f0c75.pt",48+    "dpt_large": f"{root_dir}/annotator/ckpts/dpt_large-midas-2f21e586.pt",49+    "dpt_hybrid": f"{root_dir}/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt",50     "midas_v21": "",51     "midas_v21_small": "",52 }53diff --git a/annotator/mlsd/__init__.py b/annotator/mlsd/__init__.py54index 75db717..f310fe6 10064455--- a/annotator/mlsd/__init__.py56+++ b/annotator/mlsd/__init__.py57@@ -1,3 +1,5 @@58+import pathlib59+60 import cv261 import numpy as np62 import torch63@@ -8,8 +10,9 @@ from .models.mbv2_mlsd_tiny import  MobileV2_MLSD_Tiny64 from .models.mbv2_mlsd_large import  MobileV2_MLSD_Large65 from .utils import  pred_lines66 67+root_dir = pathlib.Path(__file__).parents[2]68 69-model_path = './annotator/ckpts/mlsd_large_512_fp32.pth'70+model_path = f'{root_dir}/annotator/ckpts/mlsd_large_512_fp32.pth'71 model = MobileV2_MLSD_Large()72 model.load_state_dict(torch.load(model_path), strict=True)73 model = model.cuda().eval()74diff --git a/annotator/openpose/__init__.py b/annotator/openpose/__init__.py75index 47d50a5..2369eed 10064476--- a/annotator/openpose/__init__.py77+++ b/annotator/openpose/__init__.py78@@ -1,4 +1,5 @@79 import os80+import pathlib81 os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"82 83 import torch84@@ -7,8 +8,10 @@ from . import util85 from .body import Body86 from .hand import Hand87 88-body_estimation = Body('./annotator/ckpts/body_pose_model.pth')89-hand_estimation = Hand('./annotator/ckpts/hand_pose_model.pth')90+root_dir = pathlib.Path(__file__).parents[2]91+92+body_estimation = Body(f'{root_dir}/annotator/ckpts/body_pose_model.pth')93+hand_estimation = Hand(f'{root_dir}/annotator/ckpts/hand_pose_model.pth')94 95 96 def apply_openpose(oriImg, hand=False):97diff --git a/annotator/uniformer/__init__.py b/annotator/uniformer/__init__.py98index 500e53c..4061dbe 10064499--- a/annotator/uniformer/__init__.py100+++ b/annotator/uniformer/__init__.py101@@ -1,9 +1,12 @@102+import pathlib103+104 from annotator.uniformer.mmseg.apis import init_segmentor, inference_segmentor, show_result_pyplot105 from annotator.uniformer.mmseg.core.evaluation import get_palette106 107+root_dir = pathlib.Path(__file__).parents[2]108 109-checkpoint_file = "annotator/ckpts/upernet_global_small.pth"110-config_file = 'annotator/uniformer/exp/upernet_global_small/config.py'111+checkpoint_file = f"{root_dir}/annotator/ckpts/upernet_global_small.pth"112+config_file = f'{root_dir}/annotator/uniformer/exp/upernet_global_small/config.py'113 model = init_segmentor(config_file, checkpoint_file).cuda()114 115 116