From 553187872eeb1d52e50dc53209fa57e569609a72 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Fri, 6 Mar 2026 18:38:50 +0200 Subject: [PATCH] Fixes --- depth_anything_v2/dinov2.py | 44 ++++++------------------------------- nodes.py | 39 +++++++++++--------------------- pyproject.toml | 2 +- 3 files changed, 21 insertions(+), 64 deletions(-) diff --git a/depth_anything_v2/dinov2.py b/depth_anything_v2/dinov2.py index 83d2508..ec0dfd1 100644 --- a/depth_anything_v2/dinov2.py +++ b/depth_anything_v2/dinov2.py @@ -10,30 +10,15 @@ from functools import partial import math import logging -from typing import Sequence, Tuple, Union, Callable +from typing import Sequence, Tuple, Union import torch import torch.nn as nn -import torch.utils.checkpoint -from torch.nn.init import trunc_normal_ from .dinov2_layers import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block - logger = logging.getLogger("dinov2") - -def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module: - if not depth_first and include_root: - fn(module=module, name=name) - for child_name, child_module in module.named_children(): - child_name = ".".join((name, child_name)) if name else child_name - named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True) - if depth_first and include_root: - fn(module=module, name=name) - return module - - class BlockChunk(nn.ModuleList): def forward(self, x): for b in self: @@ -167,14 +152,7 @@ class DinoVisionTransformer(nn.Module): self.mask_token = nn.Parameter(torch.zeros(1, embed_dim)) - self.init_weights() - - def init_weights(self): - trunc_normal_(self.pos_embed, std=0.02) - nn.init.normal_(self.cls_token, std=1e-6) - if self.register_tokens is not None: - nn.init.normal_(self.register_tokens, std=1e-6) - named_apply(init_weights_vit_timm, self) + pass def interpolate_pos_encoding(self, x, w, h): previous_dtype = x.dtype @@ -193,7 +171,6 @@ class DinoVisionTransformer(nn.Module): # DINOv2 with register modify the interpolate_offset from 0.1 to 0.0 w0, h0 = w0 + self.interpolate_offset, h0 + self.interpolate_offset # w0, h0 = w0 + 0.1, h0 + 0.1 - sqrt_N = math.sqrt(N) sx, sy = float(w0) / sqrt_N, float(h0) / sqrt_N patch_pos_embed = nn.functional.interpolate( @@ -203,7 +180,7 @@ class DinoVisionTransformer(nn.Module): mode="bicubic", antialias=self.interpolate_antialias ) - + assert int(w0) == patch_pos_embed.shape[-2] assert int(h0) == patch_pos_embed.shape[-1] patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) @@ -328,13 +305,6 @@ class DinoVisionTransformer(nn.Module): return self.head(ret["x_norm_clstoken"]) -def init_weights_vit_timm(module: nn.Module, name: str = ""): - """ViT weight initialization, original timm impl (for reproducibility)""" - if isinstance(module, nn.Linear): - trunc_normal_(module.weight, std=0.02) - if module.bias is not None: - nn.init.zeros_(module.bias) - def vit_small(patch_size=16, num_register_tokens=0, **kwargs): model = DinoVisionTransformer( @@ -397,12 +367,12 @@ def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs): def DINOv2(model_name): model_zoo = { - "vits": vit_small, - "vitb": vit_base, - "vitl": vit_large, + "vits": vit_small, + "vitb": vit_base, + "vitl": vit_large, "vitg": vit_giant2 } - + return model_zoo[model_name]( img_size=518, patch_size=14, diff --git a/nodes.py b/nodes.py index 116b84f..ac5559a 100644 --- a/nodes.py +++ b/nodes.py @@ -11,15 +11,6 @@ import folder_paths from .depth_anything_v2.dpt import DepthAnythingV2 -from contextlib import nullcontext -try: - from accelerate import init_empty_weights - from accelerate.utils import set_module_tensor_to_device - is_accelerate_available = True -except: - is_accelerate_available = False - pass - class DownloadAndLoadDepthAnythingV2Model: @classmethod def INPUT_TYPES(s): @@ -51,14 +42,13 @@ class DownloadAndLoadDepthAnythingV2Model: FUNCTION = "loadmodel" CATEGORY = "DepthAnythingV2" DESCRIPTION = """ -Models autodownload to `ComfyUI\models\depthanything` from +Models autodownload to `ComfyUI/models/depthanything` from https://huggingface.co/Kijai/DepthAnythingV2-safetensors/tree/main fp16 reduces quality by a LOT, not recommended. """ def loadmodel(self, model, precision="fp32"): - device = mm.get_torch_device() if precision == "auto": dtype = torch.float16 if "fp16" in model else torch.float32 elif precision == "bf16": @@ -107,18 +97,13 @@ fp16 reduces quality by a LOT, not recommended. else: max_depth = 80.0 - with (init_empty_weights() if is_accelerate_available else nullcontext()): - if 'metric' in model: - self.model = DepthAnythingV2(**{**model_configs[encoder], 'is_metric': True, 'max_depth': max_depth}) - else: - self.model = DepthAnythingV2(**model_configs[encoder]) - - state_dict = load_torch_file(model_path) - if is_accelerate_available: - for key in state_dict: - set_module_tensor_to_device(self.model, key, device=device, dtype=dtype, value=state_dict[key]) + if 'metric' in model: + self.model = DepthAnythingV2(**{**model_configs[encoder], 'is_metric': True, 'max_depth': max_depth}) else: - self.model.load_state_dict(state_dict) + self.model = DepthAnythingV2(**model_configs[encoder]) + + state_dict = load_torch_file(model_path) + self.model.load_state_dict(state_dict, strict=False) self.model.eval() @@ -182,19 +167,21 @@ https://depth-anything-v2.github.io model.to(offload_device) mm.soft_empty_cache() depth_out = torch.cat(out, dim=0) - depth_out = depth_out.unsqueeze(-1).repeat(1, 1, 1, 3).cpu().float() + depth_out = depth_out.unsqueeze(-1).expand(-1, -1, -1, 3).cpu().float() final_H = (orig_H // 2) * 2 final_W = (orig_W // 2) * 2 if depth_out.shape[1] != final_H or depth_out.shape[2] != final_W: depth_out = F.interpolate(depth_out.permute(0, 3, 1, 2), size=(final_H, final_W), mode="bilinear").permute(0, 2, 3, 1) - depth_out = (depth_out - depth_out.min()) / (depth_out.max() - depth_out.min()) - depth_out = torch.clamp(depth_out, 0, 1) + depth_min = depth_out.min() + depth_max = depth_out.max() + depth_out.sub_(depth_min).div_(depth_max - depth_min) + depth_out.clamp_(0, 1) if da_model['is_metric']: depth_out = 1 - depth_out return (depth_out,) - + NODE_CLASS_MAPPINGS = { "DepthAnything_V2": DepthAnything_V2, "DownloadAndLoadDepthAnythingV2Model": DownloadAndLoadDepthAnythingV2Model diff --git a/pyproject.toml b/pyproject.toml index 5b98871..24ecf1c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-depthanythingv2" description = "ComfyUI nodes to use [a/DepthAnythingV2](https://depth-anything-v2.github.io/)\nNOTE:Models autodownload to ComfyUI/models/depthanything from [a/https://huggingface.co/Kijai/DepthAnythingV2-safetensors/tree/main](https://huggingface.co/Kijai/DepthAnythingV2-safetensors/tree/main)" -version = "1.0.1" +version = "1.0.2" license = "LICENSE" dependencies = ["huggingface_hub", "accelerate"]