Yuvraj SeegolamandClaude Sonnet 4.5
9ec12e4c17
- Implement ColoredLogger with lightning emoji for enhanced visual feedback - Add color-coded log levels (INFO=Green, WARNING=Yellow, ERROR=Red) - Refactor model configurations to config/models.json for easier customization - Integrate huggingface_hub for reliable downloads with progress tracking - Add tqdm progress bar styled to match logger format for batch processing - Improve ONNX model version detection for DA3, DAD, and custom models - Reorganize codebase: move nodes to nodes/ directory, utilities to separate files - Update dependencies: add huggingface_hub to requirements.txt - Enhance error handling and logging throughout engine builder and depth processing Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
58 lines
1.4 KiB
Python
58 lines
1.4 KiB
Python
import torch
|
|
import time
|
|
import sys
|
|
import os
|
|
import argparse
|
|
|
|
sys.path.append(os.path.dirname(__file__))
|
|
|
|
try:
|
|
from ..trt_utilities import Engine
|
|
except ImportError:
|
|
from trt_utilities import Engine
|
|
|
|
def export_trt(trt_path: str, onnx_path: str, use_fp16: bool):
|
|
if not os.path.isfile(onnx_path):
|
|
raise FileNotFoundError(f"Onnx file doesn't exist: {onnx_path}")
|
|
|
|
engine = Engine(trt_path)
|
|
|
|
torch.cuda.empty_cache()
|
|
|
|
s = time.time()
|
|
ret = engine.build(
|
|
onnx_path,
|
|
use_fp16,
|
|
enable_preview=True,
|
|
)
|
|
e = time.time()
|
|
print(f"Time taken to build: {(e-s)} seconds")
|
|
|
|
return ret
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(
|
|
description="Export TensorRT engine from ONNX model."
|
|
)
|
|
parser.add_argument(
|
|
"--trt-path",
|
|
type=str,
|
|
default="./depth_anything_vitl14-fp16.engine",
|
|
help="Path to save the TensorRT engine file.",
|
|
)
|
|
parser.add_argument(
|
|
"--onnx-path",
|
|
type=str,
|
|
default="./depth_anything_vitl14.onnx",
|
|
help="Path to the ONNX model file.",
|
|
)
|
|
parser.add_argument(
|
|
"--use-fp32",
|
|
action="store_true",
|
|
help="Use FP32 precision (default is FP16).",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
export_trt(
|
|
trt_path=args.trt_path, onnx_path=args.onnx_path, use_fp16=not args.use_fp32
|
|
) |