feat: add samples

README.md

@@ -10,7 +10,8 @@ Thanks & Inspired: kohya-ss/sd-webui-additional-networks
 ### Limits
 
 * Dragging large file on the Web UI may freeze the entire page. It is better to use the upload file option instead.
-* Batch size > 1 or Latent Upscale will encounter errors. (will fix later)
+* Batch size > 1 or Latent Upscale will encounter errors. 
+* MiDas Mode not working due to init issue.
 
 ### Install
 
@@ -35,6 +36,7 @@ Currently it supports both full models and trimmed models. Use `extract_controln
 | Source | Input | Output |
 |:-------------------------:|:-------------------------:|:-------------------------:|
 | (no preprocessor) |  <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/bal-source.png?raw=true"> | <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/bal-gen.png?raw=true"> |
+| (no preprocessor) |  <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/dog_rel.jpg?raw=true"> | <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/dog_rel.png?raw=true"> |
 |<img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/mahiro_input.png?raw=true">  |  <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/mahiro_canny.png?raw=true"> | <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/mahiro_gen.png?raw=true"> |
 |<img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/evt_source.jpg?raw=true">  |  <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/evt_hed.png?raw=true"> | <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/evt_gen.png?raw=true"> |
 |<img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/an-source.jpg?raw=true">  |  <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/an-pose.png?raw=true"> | <img width="256" alt="" src="https://github.com/Mikubill/sd-webui-controlnet/blob/main/samples/an-gen.png?raw=true"> |

annotator/hed/__init__.py

@@ -96,8 +96,8 @@ class Network(torch.nn.Module):
 # end
 
 netNetwork = None
-remote_model_path = "https://huggingface.co/datasets/nyanko7/tmp-public/resolve/main/network-bsds500.pt"
-modeldir = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator")
+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")
 
 def apply_hed(input_image):
     global netNetwork

annotator/mlsd/__init__.py

@@ -7,15 +7,24 @@ from einops import rearrange
 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
 
-
-model_path = './annotator/ckpts/mlsd_large_512_fp32.pth'
-model = MobileV2_MLSD_Large()
-model.load_state_dict(torch.load(model_path), strict=True)
-model = model.cuda().eval()
-
+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", "mlsd")
 
 def apply_mlsd(input_image, thr_v, thr_d):
+    global modelpath
+    if mlsdmodel is None:
+        modelpath = os.path.join(modeldir, "mlsd_large_512_fp32.pth")
+        if 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()
+        mlsdmodel.load_state_dict(torch.load(modelpath), strict=True)
+        mlsdmodel = mlsdmodel.cuda().eval()
+        
+    model = mlsdmodel
     assert input_image.ndim == 3
     img = input_image
     img_output = np.zeros_like(img)

annotator/openpose/__init__.py

@@ -13,7 +13,7 @@ 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", "openpose")
+modeldir = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "openpose")
 
 def apply_openpose(oriImg, hand=False):
     global body_estimation, hand_estimation

scripts/cldm.py

@@ -14,9 +14,7 @@ from ldm.modules.diffusionmodules.util import (
 
 from ldm.modules.attention import SpatialTransformer
 from ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample, AttentionBlock
-from ldm.util import log_txt_as_img, exists, instantiate_from_config
-from ldm.models.diffusion.ddpm import LatentDiffusion
-from modules import shared, script_callbacks
+from ldm.util import exists
 
 
 def load_state_dict(ckpt_path, location='cpu'):

scripts/controlnet.py

@@ -112,6 +112,7 @@ class Script(scripts.Script):
             "midas": midas,
             "mlsd": mlsd,
             "openpose": openpose,
+            "fake_scribble": fake_scribble,
         }
         self.input_image = None
         self.latest_model_hash = ""

scripts/processor.py

@@ -26,11 +26,20 @@ def hed(img, res=512):
     result = model_hed(img)
     return result
 
+def fake_scribble(img, res=512):
+    result = hed(img, res)
+    import cv2
+    from annotator.hed import nms
+    result = nms(result, 127, 3.0)
+    result = cv2.GaussianBlur(result, (0, 0), 3.0)
+    result[result > 10] = 255
+    result[result < 255] = 0
+    return result
 
 model_mlsd = None
 
 
-def mlsd(img, res, thr_v, thr_d):
+def mlsd(img, res=512, thr_v=0.1, thr_d=0.1):
     img = resize_image(HWC3(img), res)
     global model_mlsd
     if model_mlsd is None: