173 lines
8.0 KiB
Python
173 lines
8.0 KiB
Python
import os
|
|
import folder_paths
|
|
from .export_trt import export_trt
|
|
import urllib.request
|
|
import shutil
|
|
|
|
ENGINE_DIR = os.path.join(folder_paths.models_dir,"tensorrt", "depth-anything")
|
|
ONNX_DIR = os.path.join(folder_paths.models_dir,"onnx", "depth-anything")
|
|
|
|
# Model URLs and configurations
|
|
DEPTH_ANYTHING_MODELS = {
|
|
"v1_small": {
|
|
"url": "https://huggingface.co/yuvraj108c/Depth-Anything-Onnx/resolve/main/depth_anything_vits14.onnx",
|
|
"filename": "depth_anything_vits14.onnx",
|
|
"engine_name": "v1_depth_anything_vits14-fp16.engine"
|
|
},
|
|
"v1_base": {
|
|
"url": "https://huggingface.co/yuvraj108c/Depth-Anything-Onnx/resolve/main/depth_anything_vitb14.onnx",
|
|
"filename": "depth_anything_vitb14.onnx",
|
|
"engine_name": "v1_depth_anything_vitb14-fp16.engine"
|
|
},
|
|
"v1_large": {
|
|
"url": "https://huggingface.co/yuvraj108c/Depth-Anything-Onnx/resolve/main/depth_anything_vitl14.onnx",
|
|
"filename": "depth_anything_vitl14.onnx",
|
|
"engine_name": "v1_depth_anything_vitl14-fp16.engine"
|
|
},
|
|
"v2_small": {
|
|
"url": "https://huggingface.co/yuvraj108c/Depth-Anything-2-Onnx/resolve/main/depth_anything_v2_vits.onnx",
|
|
"filename": "depth_anything_v2_vits.onnx",
|
|
"engine_name": "v2_depth_anything_vits-fp16.engine"
|
|
},
|
|
"v2_base": {
|
|
"url": "https://huggingface.co/yuvraj108c/Depth-Anything-2-Onnx/resolve/main/depth_anything_v2_vitb.onnx",
|
|
"filename": "depth_anything_v2_vitb.onnx",
|
|
"engine_name": "v2_depth_anything_vitb-fp16.engine"
|
|
},
|
|
"v2_large": {
|
|
"url": "https://huggingface.co/yuvraj108c/Depth-Anything-2-Onnx/resolve/main/depth_anything_v2_vitl.onnx",
|
|
"filename": "depth_anything_v2_vitl.onnx",
|
|
"engine_name": "v2_depth_anything_vitl-fp16.engine"
|
|
},
|
|
"v2_metric_hypersim_large": {
|
|
"url": "https://huggingface.co/yuvraj108c/Depth-Anything-2-Onnx/resolve/main/depth_anything_v2_metric_hypersim_vitl.onnx",
|
|
"filename": "depth_anything_v2_metric_hypersim_vitl.onnx",
|
|
"engine_name": "v2_depth_anything_metric_hypersim_vitl-fp16.engine"
|
|
},
|
|
"v2_metric_vkitti_large": {
|
|
"url": "https://huggingface.co/yuvraj108c/Depth-Anything-2-Onnx/resolve/main/depth_anything_v2_metric_vkitti_vitl.onnx",
|
|
"filename": "depth_anything_v2_metric_vkitti_vitl.onnx",
|
|
"engine_name": "v2_depth_anything_metric_vkitti_vitl-fp16.engine"
|
|
},
|
|
"DAD_small": {
|
|
"url": "https://huggingface.co/yuvraj108c/distill-any-depth-onnx/resolve/main/Distill-Any-Depth-Multi-Teacher-Small.onnx",
|
|
"filename": "Distill-Any-Depth-Multi-Teacher-Small.onnx",
|
|
"engine_name": "Distill-Any-Depth-Multi-Teacher-Small-fp16.engine"
|
|
},
|
|
"DAD_base": {
|
|
"url": "https://huggingface.co/yuvraj108c/distill-any-depth-onnx/resolve/main/Distill-Any-Depth-Multi-Teacher-Base.onnx",
|
|
"filename": "Distill-Any-Depth-Multi-Teacher-Base.onnx",
|
|
"engine_name": "Distill-Any-Depth-Multi-Teacher-Base-fp16.engine"
|
|
},
|
|
"DAD_large": {
|
|
"url": "https://huggingface.co/yuvraj108c/distill-any-depth-onnx/resolve/main/Distill-Any-Depth-Multi-Teacher-Large.onnx",
|
|
"filename": "Distill-Any-Depth-Multi-Teacher-Large.onnx",
|
|
"engine_name": "Distill-Any-Depth-Multi-Teacher-Large-fp16.engine"
|
|
},
|
|
"DAD_large_2w_iter": {
|
|
"url": "https://huggingface.co/yuvraj108c/distill-any-depth-onnx/resolve/main/Distill-Any-Depth-Dav2-Teacher-Large-2w-iter.onnx",
|
|
"filename": "Distill-Any-Depth-Dav2-Teacher-Large-2w-iter.onnx",
|
|
"engine_name": "Distill-Any-Depth-Dav2-Teacher-Large-2w-iter-fp16.engine"
|
|
}
|
|
}
|
|
|
|
def download_onnx_model(model_key):
|
|
"""Download ONNX model if it doesn't exist"""
|
|
if model_key not in DEPTH_ANYTHING_MODELS:
|
|
return f"Invalid model key: {model_key}", None
|
|
|
|
model_info = DEPTH_ANYTHING_MODELS[model_key]
|
|
url = model_info["url"]
|
|
filename = model_info["filename"]
|
|
local_path = os.path.join(ONNX_DIR, filename)
|
|
|
|
if os.path.exists(local_path):
|
|
return f"Model already exists at: {local_path}", local_path
|
|
|
|
try:
|
|
print(f"Downloading model from {url} to {local_path}...")
|
|
|
|
# Create a temporary file for downloading
|
|
temp_file = local_path + ".tmp"
|
|
|
|
# Download the file
|
|
with urllib.request.urlopen(url) as response, open(temp_file, 'wb') as out_file:
|
|
shutil.copyfileobj(response, out_file)
|
|
|
|
# Rename the temporary file to the target filename
|
|
shutil.move(temp_file, local_path)
|
|
|
|
return f"Successfully downloaded model to: {local_path}", local_path
|
|
except Exception as e:
|
|
return f"Error downloading model: {str(e)}", None
|
|
|
|
class DepthAnythingEngineBuilder:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_version": (list(DEPTH_ANYTHING_MODELS.keys()), {
|
|
"default": "v1_large",
|
|
"tooltip": "Select the Depth Anything model version and size. Larger models provide better quality but require more VRAM."
|
|
}),
|
|
"custom_engine_name": ("STRING", {
|
|
"default": "",
|
|
"tooltip": "Optional custom name for the TensorRT engine file. If empty, will use the default name based on the model."
|
|
}),
|
|
"use_fp16": ("BOOLEAN", {
|
|
"default": True,
|
|
"tooltip": "Enable FP16 precision for faster inference and lower VRAM usage. Disable if you experience stability issues."
|
|
}),
|
|
"custom_onnx_path": ("STRING", {
|
|
"default": "",
|
|
"optional": True,
|
|
"tooltip": "Optional path to a custom ONNX model file. If provided, will use this instead of downloading the predefined model."
|
|
}),
|
|
}
|
|
}
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("message",)
|
|
FUNCTION = "build_engine"
|
|
CATEGORY = "tensorrt"
|
|
OUTPUT_NODE = True
|
|
|
|
def build_engine(self, model_version, custom_engine_name, use_fp16, custom_onnx_path=""):
|
|
# Ensure directories exist
|
|
os.makedirs(ENGINE_DIR, exist_ok=True)
|
|
os.makedirs(ONNX_DIR, exist_ok=True)
|
|
|
|
# Determine ONNX path and engine name
|
|
if custom_onnx_path and os.path.exists(custom_onnx_path):
|
|
onnx_path = custom_onnx_path
|
|
version = "v1" if "v1" in os.path.basename(custom_onnx_path) else "v2"
|
|
engine_name = custom_engine_name if custom_engine_name else f"{version}_{os.path.basename(custom_onnx_path).replace('.onnx', '-fp16.engine' if use_fp16 else '.engine')}"
|
|
else:
|
|
# Download and use predefined model
|
|
model_info = DEPTH_ANYTHING_MODELS[model_version]
|
|
message, onnx_path = download_onnx_model(model_version)
|
|
|
|
if not onnx_path:
|
|
return (message,)
|
|
|
|
engine_name = custom_engine_name if custom_engine_name else model_info["engine_name"]
|
|
if use_fp16 and not engine_name.endswith("-fp16.engine"):
|
|
engine_name = engine_name.replace(".engine", "-fp16.engine")
|
|
elif not use_fp16 and "-fp16.engine" in engine_name:
|
|
engine_name = engine_name.replace("-fp16.engine", ".engine")
|
|
|
|
engine_path = os.path.join(ENGINE_DIR, engine_name)
|
|
|
|
# Check if engine already exists
|
|
if os.path.exists(engine_path):
|
|
return (f"Engine already exists at: {engine_path}. Delete it manually if you want to rebuild.",)
|
|
|
|
try:
|
|
print(f"Building TensorRT engine from {onnx_path} to {engine_path}...")
|
|
result = export_trt(trt_path=engine_path, onnx_path=onnx_path, use_fp16=use_fp16)
|
|
if result == 0:
|
|
return (f"Successfully built engine: {engine_path}",)
|
|
else:
|
|
return (f"Failed to build engine. Check console for details.",)
|
|
except Exception as e:
|
|
return (f"Error building engine: {str(e)}",)
|