Merge pull request AUTOMATIC1111#138 from missionfloyd/model-path

Move models to models/ControlNet

README.md

@@ -26,7 +26,7 @@ Upgrade gradio if any ui issues occured: `pip install gradio==3.16.2`
 
 ### Usage
 
-1. Put the ControlNet models (`.pt`, `.pth`, `.ckpt` or `.safetensors`) inside the `sd-webui-controlnet/models` folder.
+1. Put the ControlNet models (`.pt`, `.pth`, `.ckpt` or `.safetensors`) inside the `models/ControlNet` folder.
 2. Open "txt2img" or "img2img" tab, write your prompts.
 3. Press "Refresh models" and select the model you want to use. (If nothing appears, try reload/restart the webui)
 4. Upload your image and select preprocessor, done.

annotator/hed/__init__.py

@@ -1,10 +1,12 @@
+from distutils import extension
 import numpy as np
 import cv2
 import torch
 from einops import rearrange
 
 import os 
-from modules import extensions, devices
+from modules import devices
+from modules.paths import models_path
 
 class Network(torch.nn.Module):
     def __init__(self, model_path):
@@ -97,13 +99,17 @@ def forward(self, tenInput):
 
 netNetwork = None
 remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/network-bsds500.pth"
-modeldir = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "hed")
+modeldir = os.path.join(models_path, "hed")
+old_modeldir = os.path.dirname(os.path.realpath(__file__))
 
 def apply_hed(input_image):
     global netNetwork
     if netNetwork is None:
         modelpath = os.path.join(modeldir, "network-bsds500.pth")
-        if not os.path.exists(modelpath):
+        old_modelpath = os.path.join(old_modeldir, "network-bsds500.pth")
+        if os.path.exists(old_modelpath):
+            modelpath = old_modelpath
+        elif not os.path.exists(modelpath):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(remote_model_path, model_dir=modeldir)
         netNetwork = Network(modelpath)

annotator/leres/__init__.py

@@ -2,7 +2,8 @@
 import numpy as np
 import torch
 import os
-from modules import extensions, devices, shared
+from modules import devices, shared
+from modules.paths import models_path
 from torchvision.transforms import transforms
 
 # AdelaiDepth/LeReS imports
@@ -14,7 +15,9 @@
 from .pix2pix.options.test_options import TestOptions
 from .pix2pix.models.pix2pix4depth_model import Pix2Pix4DepthModel
 
-base_model_path = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "leres")
+base_model_path = os.path.join(models_path, "leres")
+old_modeldir = os.path.dirname(os.path.realpath(__file__))
+
 remote_model_path_leres = "https://cloudstor.aarnet.edu.au/plus/s/lTIJF4vrvHCAI31/download"
 remote_model_path_pix2pix = "https://sfu.ca/~yagiz/CVPR21/latest_net_G.pth"
 
@@ -34,7 +37,11 @@ def apply_leres(input_image, thr_a, thr_b):
 
     if model is None:
         model_path = os.path.join(base_model_path, "res101.pth")
-        if not os.path.exists(model_path):
+        old_model_path = os.path.join(old_modeldir, "res101.pth")
+
+        if os.path.exists(old_model_path):
+            model_path = old_model_path
+        elif not os.path.exists(model_path):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(remote_model_path_leres, model_dir=base_model_path)
             os.rename(os.path.join(base_model_path, 'download'), model_path)

annotator/midas/api.py

@@ -4,7 +4,7 @@
 import torch
 import torch.nn as nn
 import os
-from modules import extensions
+from modules.paths import models_path
 
 from torchvision.transforms import Compose
 
@@ -13,7 +13,8 @@
 from .midas.midas_net_custom import MidasNet_small
 from .midas.transforms import Resize, NormalizeImage, PrepareForNet
 
-base_model_path = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "midas")
+base_model_path = os.path.join(models_path, "midas")
+old_modeldir = os.path.dirname(os.path.realpath(__file__))
 remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/dpt_hybrid-midas-501f0c75.pt"
 
 ISL_PATHS = {
@@ -23,6 +24,13 @@
     "midas_v21_small": "",
 }
 
+OLD_ISL_PATHS = {
+    "dpt_large": os.path.join(old_modeldir, "dpt_large-midas-2f21e586.pt"),
+    "dpt_hybrid": os.path.join(old_modeldir, "dpt_hybrid-midas-501f0c75.pt"),
+    "midas_v21": "",
+    "midas_v21_small": "",
+}
+
 
 def disabled_train(self, mode=True):
     """Overwrite model.train with this function to make sure train/eval mode
@@ -79,6 +87,7 @@ def load_model(model_type):
     # https://github.com/isl-org/MiDaS/blob/master/run.py
     # load network
     model_path = ISL_PATHS[model_type]
+    old_model_path = OLD_ISL_PATHS[model_type]
     if model_type == "dpt_large":  # DPT-Large
         model = DPTDepthModel(
             path=model_path,
@@ -90,7 +99,9 @@ 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(model_path):
+        if os.path.exists(old_model_path):
+            model_path = old_model_path
+        elif not os.path.exists(model_path):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(remote_model_path, model_dir=base_model_path)
 

annotator/mlsd/__init__.py

@@ -7,11 +7,13 @@
 from .models.mbv2_mlsd_tiny import  MobileV2_MLSD_Tiny
 from .models.mbv2_mlsd_large import  MobileV2_MLSD_Large
 from .utils import  pred_lines
-from modules import extensions, devices
+from modules import devices
+from modules.paths import models_path
 
 mlsdmodel = None
 remote_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/mlsd_large_512_fp32.pth"
-modeldir = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "mlsd")
+old_modeldir = os.path.dirname(os.path.realpath(__file__))
+modeldir = os.path.join(models_path, "mlsd")
 
 def unload_mlsd_model():
     global mlsdmodel
@@ -22,7 +24,10 @@ def apply_mlsd(input_image, thr_v, thr_d):
     global modelpath, mlsdmodel
     if mlsdmodel is None:
         modelpath = os.path.join(modeldir, "mlsd_large_512_fp32.pth")
-        if not os.path.exists(modelpath):
+        old_modelpath = os.path.join(old_modeldir, "mlsd_large_512_fp32.pth")
+        if os.path.exists(old_modelpath):
+            modelpath = old_modelpath
+        elif not os.path.exists(modelpath):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(remote_model_path, model_dir=modeldir)
         mlsdmodel = MobileV2_MLSD_Large()

annotator/openpose/__init__.py

@@ -6,14 +6,15 @@
 from . import util
 from .body import Body
 from .hand import Hand
-from modules import extensions
+from modules.paths import models_path
 
 body_estimation = None 
 hand_estimation = None 
 
 body_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/body_pose_model.pth"
 hand_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/hand_pose_model.pth"
-modeldir = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "openpose")
+modeldir = os.path.join(models_path, "openpose")
+old_modeldir = os.path.dirname(os.path.realpath(__file__))
 
 def unload_openpose_model():
     global body_estimation, hand_estimation
@@ -26,12 +27,18 @@ def apply_openpose(oriImg, hand=False):
     if body_estimation is None:
         body_modelpath = os.path.join(modeldir, "body_pose_model.pth")
         hand_modelpath = os.path.join(modeldir, "hand_pose_model.pth")
+        old_body_modelpath = os.path.join(old_modeldir, "body_pose_model.pth")
+        old_hand_modelpath = os.path.join(old_modeldir, "hand_pose_model.pth")
 
-        if not os.path.exists(hand_modelpath):
+        if os.path.exists(old_body_modelpath):
+            body_modelpath = old_body_modelpath
+        elif not os.path.exists(hand_modelpath):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(hand_model_path, model_dir=modeldir)
 
-        if not os.path.exists(body_model_path):
+        if os.path.exists(old_hand_modelpath):
+            hand_modelpath = old_hand_modelpath
+        elif not os.path.exists(body_model_path):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(body_model_path, model_dir=modeldir)
 

annotator/openpose/body.py

@@ -1,3 +1,4 @@
+import os
 import cv2
 import numpy as np
 import math
@@ -11,6 +12,7 @@
 from . import util
 from .model import bodypose_model
 from modules import devices
+from modules.paths import models_path
 
 class Body(object):
     def __init__(self, model_path):
@@ -209,7 +211,7 @@ def __call__(self, oriImg):
         return candidate, subset
 
 if __name__ == "__main__":
-    body_estimation = Body('../model/body_pose_model.pth')
+    body_estimation = Body(os.path.join(models_path, "openpose", "body_pose_model.pth"))
 
     test_image = '../images/ski.jpg'
     oriImg = cv2.imread(test_image)  # B,G,R order

annotator/openpose/hand.py

@@ -12,6 +12,7 @@
 from .model import handpose_model
 from . import util
 from modules import devices
+from modules.paths import models_path
 
 class Hand(object):
     def __init__(self, model_path):
@@ -75,7 +76,7 @@ def __call__(self, oriImg):
         return np.array(all_peaks)
 
 if __name__ == "__main__":
-    hand_estimation = Hand('../model/hand_pose_model.pth')
+    hand_estimation = Hand(os.path.join(models_path, "openpose", "hand_pose_model.pth"))
 
     # test_image = '../images/hand.jpg'
     test_image = '../images/hand.jpg'

annotator/pidinet/__init__.py

@@ -3,18 +3,23 @@
 import numpy as np
 from einops import rearrange
 from annotator.pidinet.model import pidinet
-from modules import extensions, devices
+from modules import devices
+from modules.paths import models_path
 from scripts.utils import load_state_dict
 
 netNetwork = None
 remote_model_path = "https://github.com/TencentARC/T2I-Adapter/raw/main/models/table5_pidinet.pth"
-modeldir = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "pidinet")
+modeldir = os.path.join(models_path, "pidinet")
+old_modeldir = os.path.dirname(os.path.realpath(__file__))
 
 def apply_pidinet(input_image):
     global netNetwork
     if netNetwork is None:
         modelpath = os.path.join(modeldir, "table5_pidinet.pth")
-        if not os.path.exists(modelpath):
+        old_modelpath = os.path.join(old_modeldir, "table5_pidinet.pth")
+        if os.path.exists(old_modelpath):
+            modelpath = old_modelpath
+        elif not os.path.exists(modelpath):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(remote_model_path, model_dir=modeldir)
         netNetwork = pidinet()

annotator/uniformer/__init__.py

@@ -1,12 +1,14 @@
 from annotator.uniformer.mmseg.apis import init_segmentor, inference_segmentor, show_result_pyplot
 from annotator.uniformer.mmseg.core.evaluation import get_palette
 from modules.shared import extensions
+from modules.paths import models_path
 from modules import devices
 import os
 
-modeldir = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "uniformer")
+modeldir = os.path.join(models_path, "uniformer")
 checkpoint_file = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/upernet_global_small.pth"
-config_file = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "uniformer", "exp", "upernet_global_small", "config.py")
+config_file = os.path.join(os.path.dirname(os.path.realpath(__file__)), "exp", "upernet_global_small", "config.py")
+old_modeldir = os.path.dirname(os.path.realpath(__file__))
 model = None
 
 def unload_uniformer_model():
@@ -18,7 +20,10 @@ def apply_uniformer(img):
     global model
     if model is None:
         modelpath = os.path.join(modeldir, "upernet_global_small.pth")
-        if not os.path.exists(modelpath):
+        old_modelpath = os.path.join(old_modeldir, "upernet_global_small.pth")
+        if os.path.exists(old_modelpath):
+            modelpath = old_modelpath  
+        elif not os.path.exists(modelpath):
             from basicsr.utils.download_util import load_file_from_url
             load_file_from_url(checkpoint_file, model_dir=modeldir)
 

scripts/controlnet.py

@@ -17,6 +17,7 @@
 from scripts.utils import load_state_dict
 from scripts.hook import ControlParams, UnetHook
 from modules import sd_models
+from modules.paths import models_path
 from modules.processing import StableDiffusionProcessingImg2Img
 from modules.images import save_image
 from PIL import Image
@@ -45,12 +46,13 @@
 CN_MODEL_EXTS = [".pt", ".pth", ".ckpt", ".safetensors"]
 cn_models = {}      # "My_Lora(abcd1234)" -> C:/path/to/model.safetensors
 cn_models_names = {}  # "my_lora" -> "My_Lora(abcd1234)"
-cn_models_dir = os.path.join(scripts.basedir(), "models")
+cn_models_dir = os.path.join(models_path, "ControlNet")
+cn_models_dir_old = os.path.join(scripts.basedir(), "models")
 os.makedirs(cn_models_dir, exist_ok=True)
+default_conf = os.path.join(scripts.basedir(), "models", "cldm_v15.yaml")
+default_conf_adapter = os.path.join(scripts.basedir(), "models", "sketch_adapter_v14.yaml")
 cn_detectedmap_dir = os.path.join(scripts.basedir(), "detected_maps")
 os.makedirs(cn_detectedmap_dir, exist_ok=True)
-default_conf_adapter = os.path.join(cn_models_dir, "sketch_adapter_v14.yaml")
-default_conf = os.path.join(cn_models_dir, "cldm_v15.yaml")
 default_detectedmap_dir = cn_detectedmap_dir
 refresh_symbol = '\U0001f504'       # 🔄
 switch_values_symbol = '\U000021C5' # ⇅
@@ -138,7 +140,7 @@ def update_cn_models():
     ext_dirs = (shared.opts.data.get("control_net_models_path", None), getattr(shared.cmd_opts, 'controlnet_dir', None))
     extra_lora_paths = (extra_lora_path for extra_lora_path in ext_dirs
                 if extra_lora_path is not None and os.path.exists(extra_lora_path))
-    paths = [cn_models_dir, *extra_lora_paths]
+    paths = [cn_models_dir, cn_models_dir_old, *extra_lora_paths]
 
     for path in paths:
         sort_by = shared.opts.data.get(