Adding Depth Anything v2 to ControlNet Integrated (#1519)

* Update install.py

adding install of depth anything v2

* Add files via upload

adding depth anything v2 preprocessor

* Update preprocessor_compiled.py

adding preprocessor

* Update preprocessor.py

adding preprocessor functions

extensions-builtin/forge_legacy_preprocessors/annotator/depth_anything_v2.py

@@ -0,0 +1,78 @@
+import os
+import torch
+import cv2
+import numpy as np
+import torch.nn.functional as F
+from torchvision.transforms import Compose
+from safetensors.torch import load_file
+
+from depth_anything_v2.dpt import DepthAnythingV2
+from depth_anything_v2.util.transform import Resize, NormalizeImage, PrepareForNet
+from .util import load_model
+from .annotator_path import models_path
+
+transform = Compose(
+    [
+        Resize(
+            width=518,
+            height=518,
+            resize_target=False,
+            keep_aspect_ratio=True,
+            ensure_multiple_of=14,
+            resize_method="lower_bound",
+            image_interpolation_method=cv2.INTER_CUBIC,
+        ),
+        NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        PrepareForNet(),
+    ]
+)
+
+class DepthAnythingV2Detector:
+    """https://github.com/MackinationsAi/Upgraded-Depth-Anything-V2"""
+
+    model_dir = os.path.join(models_path, "depth_anything_v2")
+
+    def __init__(self, device: torch.device):
+        self.device = device
+        self.model = (
+            DepthAnythingV2(
+                encoder="vitl",
+                features=256,
+                out_channels=[256, 512, 1024, 1024],
+            )
+            .to(device)
+            .eval()
+        )
+        remote_url = os.environ.get(
+            "CONTROLNET_DEPTH_ANYTHING_V2_MODEL_URL",
+            "https://huggingface.co/MackinationsAi/Depth-Anything-V2_Safetensors/resolve/main/depth_anything_v2_vitl.safetensors",
+        )
+        model_path = load_model(
+            "depth_anything_v2_vitl.safetensors", remote_url=remote_url, model_dir=self.model_dir
+        )
+        self.model.load_state_dict(load_file(model_path))
+
+    def __call__(self, image: np.ndarray, colored: bool = True) -> np.ndarray:
+        self.model.to(self.device)
+        h, w = image.shape[:2]
+
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) / 255.0
+        image = transform({"image": image})["image"]
+        image = torch.from_numpy(image).unsqueeze(0).to(self.device)
+        @torch.no_grad()
+        def predict_depth(model, image):
+            return model(image)
+        depth = predict_depth(self.model, image)
+        depth = F.interpolate(
+            depth[None], (h, w), mode="bilinear", align_corners=False
+        )[0, 0]
+        depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
+        depth = depth.cpu().numpy().astype(np.uint8)
+        if colored:
+            depth_color = cv2.applyColorMap(depth, cv2.COLORMAP_INFERNO)[:, :, ::-1]
+            return depth_color
+        else:
+            return depth
+
+    def unload_model(self):
+        self.model.to("cpu")

extensions-builtin/forge_legacy_preprocessors/install.py

@@ -148,4 +148,13 @@ try_install_from_wheel(
         "https://github.com/huchenlei/Depth-Anything/releases/download/v1.0.0/depth_anything-2024.1.22.0-py2.py3-none-any.whl",
     ),
 )
+
+try_install_from_wheel(
+    "depth_anything_v2",
+    wheel_url=os.environ.get(
+        "DEPTH_ANYTHING_V2_WHEEL",
+        "https://github.com/MackinationsAi/UDAV2-ControlNet/releases/download/v1.0.0/depth_anything_v2-2024.7.1.0-py2.py3-none-any.whl",
+    ),
+)
+
 try_remove_legacy_submodule()

extensions-builtin/forge_legacy_preprocessors/legacy_preprocessors/preprocessor.py

@@ -208,6 +208,23 @@ def unload_depth_anything():
     if model_depth_anything is not None:
         model_depth_anything.unload_model()
 
+model_depth_anything_v2 = None
+
+
+def depth_anything_v2(img, res:int = 512, colored:bool = True, **kwargs):
+    img, remove_pad = resize_image_with_pad(img, res)
+    global model_depth_anything_v2
+    if model_depth_anything_v2 is None:
+        with Extra(torch_handler):
+            from annotator.depth_anything_v2 import DepthAnythingV2Detector
+            device = devices.get_device_for("controlnet")
+            model_depth_anything_v2 = DepthAnythingV2Detector(device)
+    return remove_pad(model_depth_anything_v2(img, colored=colored)), True
+
+
+def unload_depth_anything_v2():
+    if model_depth_anything_v2 is not None:
+        model_depth_anything_v2.unload_model()
 
 model_midas = None
 

extensions-builtin/forge_legacy_preprocessors/legacy_preprocessors/preprocessor_compiled.py

@@ -168,6 +168,22 @@ legacy_preprocessors = {
             "Depth"
         ]
     },
+   "depth_anything_v2": {
+        "label": "depth_anything_v2",
+        "call_function": functools.partial(depth_anything_v2, colored=False),
+        "unload_function": unload_depth_anything_v2,
+        "managed_model": "model_depth_anything_v2",
+        "model_free": False,
+        "no_control_mode": False,
+        "resolution": None,
+        "slider_1": None,
+        "slider_2": None,
+        "slider_3": None,
+        "priority": 0,
+        "tags": [
+            "Depth"
+        ]
+    },
     "depth_hand_refiner": {
         "label": "depth_hand_refiner",
         "call_function": g_hand_refiner_model.run_model,