From 2498b3b4672fa6ad9391903788b596db1882f8fe Mon Sep 17 00:00:00 2001 From: Varshith Bathini Date: Sun, 14 Dec 2025 19:11:55 +0000 Subject: [PATCH 1/2] fix: build time --- .../acceleration/tensorrt/builder.py | 2 + .../acceleration/tensorrt/engine_manager.py | 7 ++ .../export_wrappers/unet_controlnet_export.py | 1 - .../acceleration/tensorrt/utilities.py | 9 ++- .../tools/compile_depth_tensorrt.py | 80 +++++++++++++++++++ .../tools/compile_raft_tensorrt.py | 1 + 6 files changed, 98 insertions(+), 2 deletions(-) create mode 100644 src/streamdiffusion/tools/compile_depth_tensorrt.py diff --git a/src/streamdiffusion/acceleration/tensorrt/builder.py b/src/streamdiffusion/acceleration/tensorrt/builder.py index 8550f0c9..3843bb9c 100644 --- a/src/streamdiffusion/acceleration/tensorrt/builder.py +++ b/src/streamdiffusion/acceleration/tensorrt/builder.py @@ -46,6 +46,7 @@ def build( force_engine_build: bool = False, force_onnx_export: bool = False, force_onnx_optimize: bool = False, + timing_cache: str = None, ): if not force_onnx_export and os.path.exists(onnx_path): print(f"Found cached model: {onnx_path}") @@ -89,6 +90,7 @@ def build( build_dynamic_shape=build_dynamic_shape, build_all_tactics=build_all_tactics, build_enable_refit=build_enable_refit, + timing_cache=timing_cache, ) for file in os.listdir(os.path.dirname(engine_path)): if file.endswith('.engine'): diff --git a/src/streamdiffusion/acceleration/tensorrt/engine_manager.py b/src/streamdiffusion/acceleration/tensorrt/engine_manager.py index f5cfc1bd..f2ca141e 100644 --- a/src/streamdiffusion/acceleration/tensorrt/engine_manager.py +++ b/src/streamdiffusion/acceleration/tensorrt/engine_manager.py @@ -221,6 +221,13 @@ def compile_and_load_engine(self, Moves compilation blocks from wrapper.py lines 1200-1252, 1254-1283, 1285-1313. """ + if 'engine_build_options' not in kwargs: + kwargs['engine_build_options'] = {} + + if 'timing_cache' not in kwargs['engine_build_options']: + timing_cache_path = self.engine_dir / "timing_cache" + kwargs['engine_build_options']['timing_cache'] = str(timing_cache_path) + if not engine_path.exists(): # Get the appropriate compile function for this engine type config = self._configs[engine_type] diff --git a/src/streamdiffusion/acceleration/tensorrt/export_wrappers/unet_controlnet_export.py b/src/streamdiffusion/acceleration/tensorrt/export_wrappers/unet_controlnet_export.py index fe26a88d..d5d1f7b7 100644 --- a/src/streamdiffusion/acceleration/tensorrt/export_wrappers/unet_controlnet_export.py +++ b/src/streamdiffusion/acceleration/tensorrt/export_wrappers/unet_controlnet_export.py @@ -289,7 +289,6 @@ def forward(self, sample, timestep, encoder_hidden_states, *args): return res else: return res[0] - return res def create_controlnet_wrapper(unet: UNet2DConditionModel, diff --git a/src/streamdiffusion/acceleration/tensorrt/utilities.py b/src/streamdiffusion/acceleration/tensorrt/utilities.py index b9bfa1c3..1e5be108 100644 --- a/src/streamdiffusion/acceleration/tensorrt/utilities.py +++ b/src/streamdiffusion/acceleration/tensorrt/utilities.py @@ -19,6 +19,7 @@ # import gc +import os from collections import OrderedDict from pathlib import Path from typing import Any, List, Optional, Tuple, Union @@ -249,10 +250,14 @@ def build( if not enable_all_tactics: config_kwargs["tactic_sources"] = [] + load_timing_cache = timing_cache + if isinstance(timing_cache, (str, Path)) and not os.path.exists(timing_cache): + load_timing_cache = None + engine = engine_from_network( network_from_onnx_path(onnx_path, flags=[trt.OnnxParserFlag.NATIVE_INSTANCENORM]), config=CreateConfig( - fp16=fp16, refittable=enable_refit, profiles=[p], load_timing_cache=timing_cache, **config_kwargs + fp16=fp16, refittable=enable_refit, profiles=[p], load_timing_cache=load_timing_cache, **config_kwargs ), save_timing_cache=timing_cache, ) @@ -498,6 +503,7 @@ def build_engine( build_dynamic_shape: bool = False, build_all_tactics: bool = False, build_enable_refit: bool = False, + timing_cache: str = None, ): _, free_mem, _ = cudart.cudaMemGetInfo() GiB = 2**30 @@ -520,6 +526,7 @@ def build_engine( input_profile=input_profile, enable_refit=build_enable_refit, enable_all_tactics=build_all_tactics, + timing_cache=timing_cache, workspace_size=max_workspace_size, ) diff --git a/src/streamdiffusion/tools/compile_depth_tensorrt.py b/src/streamdiffusion/tools/compile_depth_tensorrt.py new file mode 100644 index 00000000..d7898803 --- /dev/null +++ b/src/streamdiffusion/tools/compile_depth_tensorrt.py @@ -0,0 +1,80 @@ +import os +import sys +import shutil +import logging +import subprocess +from pathlib import Path +from typing import Optional +import fire +from huggingface_hub import hf_hub_download + +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +logger = logging.getLogger(__name__) + +REPO_URL = "https://github.com/yuvraj108c/ComfyUI-Depth-Anything-Tensorrt.git" +REPO_COMMIT = "1f4c161949b3616516745781fb91444e6443cc25" +MODEL_REPO_ID = "yuvraj108c/Depth-Anything-2-Onnx" +MODEL_FILENAME = "depth_anything_v2_vits.onnx" + +def compile_depth( + output_dir: str = "engines/depth-anything", + workspace_dir: str = "workspace", + force_rebuild: bool = False, +): + """ + Compile Depth Anything TensorRT engine. + + Args: + output_dir: Directory to save the engine. + workspace_dir: Directory to clone the repository. + force_rebuild: Force rebuild even if engine exists. + """ + output_path = Path(output_dir) + workspace_path = Path(workspace_dir) + engine_path = output_path / "depth_anything_v2_vits.engine" + + if engine_path.exists() and not force_rebuild: + logger.info(f"Engine already exists: {engine_path}") + return + + # 1. Setup Workspace + repo_dir = workspace_path / "ComfyUI-Depth-Anything-Tensorrt" + if not repo_dir.exists(): + logger.info(f"Cloning {REPO_URL}...") + subprocess.run(["git", "clone", REPO_URL, str(repo_dir)], check=True) + + logger.info(f"Checking out commit {REPO_COMMIT}...") + subprocess.run(["git", "-C", str(repo_dir), "checkout", REPO_COMMIT], check=True) + + # 2. Download ONNX + logger.info(f"Downloading {MODEL_FILENAME} from {MODEL_REPO_ID}...") + onnx_path = hf_hub_download( + repo_id=MODEL_REPO_ID, + filename=MODEL_FILENAME, + ) + + # 3. Build Engine + output_path.mkdir(parents=True, exist_ok=True) + logger.info("Building TensorRT engine...") + + # Install requirements if needed + if (repo_dir / "requirements.txt").exists(): + subprocess.run([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], cwd=repo_dir, check=True) + + subprocess.run( + [ + sys.executable, + "export_trt.py", + "--trt-path", + str(engine_path.absolute()), + "--onnx-path", + str(onnx_path), + ], + cwd=repo_dir, + check=True, + ) + + logger.info(f"Successfully built engine: {engine_path}") + +if __name__ == "__main__": + fire.Fire(compile_depth) diff --git a/src/streamdiffusion/tools/compile_raft_tensorrt.py b/src/streamdiffusion/tools/compile_raft_tensorrt.py index 8dec3a76..3a5716b2 100644 --- a/src/streamdiffusion/tools/compile_raft_tensorrt.py +++ b/src/streamdiffusion/tools/compile_raft_tensorrt.py @@ -87,6 +87,7 @@ def export_raft_to_onnx( opset_version=17, export_params=True, dynamic_axes=dynamic_axes, + dynamo=False, ) del raft_model From ac74625b2c18367056e2dd562a2daf872ca15873 Mon Sep 17 00:00:00 2001 From: Varshith Bathini Date: Sun, 14 Dec 2025 19:12:44 +0000 Subject: [PATCH 2/2] fix: depth compile --- .../tools/compile_depth_tensorrt.py | 80 ------------------- 1 file changed, 80 deletions(-) delete mode 100644 src/streamdiffusion/tools/compile_depth_tensorrt.py diff --git a/src/streamdiffusion/tools/compile_depth_tensorrt.py b/src/streamdiffusion/tools/compile_depth_tensorrt.py deleted file mode 100644 index d7898803..00000000 --- a/src/streamdiffusion/tools/compile_depth_tensorrt.py +++ /dev/null @@ -1,80 +0,0 @@ -import os -import sys -import shutil -import logging -import subprocess -from pathlib import Path -from typing import Optional -import fire -from huggingface_hub import hf_hub_download - -logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') -logger = logging.getLogger(__name__) - -REPO_URL = "https://github.com/yuvraj108c/ComfyUI-Depth-Anything-Tensorrt.git" -REPO_COMMIT = "1f4c161949b3616516745781fb91444e6443cc25" -MODEL_REPO_ID = "yuvraj108c/Depth-Anything-2-Onnx" -MODEL_FILENAME = "depth_anything_v2_vits.onnx" - -def compile_depth( - output_dir: str = "engines/depth-anything", - workspace_dir: str = "workspace", - force_rebuild: bool = False, -): - """ - Compile Depth Anything TensorRT engine. - - Args: - output_dir: Directory to save the engine. - workspace_dir: Directory to clone the repository. - force_rebuild: Force rebuild even if engine exists. - """ - output_path = Path(output_dir) - workspace_path = Path(workspace_dir) - engine_path = output_path / "depth_anything_v2_vits.engine" - - if engine_path.exists() and not force_rebuild: - logger.info(f"Engine already exists: {engine_path}") - return - - # 1. Setup Workspace - repo_dir = workspace_path / "ComfyUI-Depth-Anything-Tensorrt" - if not repo_dir.exists(): - logger.info(f"Cloning {REPO_URL}...") - subprocess.run(["git", "clone", REPO_URL, str(repo_dir)], check=True) - - logger.info(f"Checking out commit {REPO_COMMIT}...") - subprocess.run(["git", "-C", str(repo_dir), "checkout", REPO_COMMIT], check=True) - - # 2. Download ONNX - logger.info(f"Downloading {MODEL_FILENAME} from {MODEL_REPO_ID}...") - onnx_path = hf_hub_download( - repo_id=MODEL_REPO_ID, - filename=MODEL_FILENAME, - ) - - # 3. Build Engine - output_path.mkdir(parents=True, exist_ok=True) - logger.info("Building TensorRT engine...") - - # Install requirements if needed - if (repo_dir / "requirements.txt").exists(): - subprocess.run([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], cwd=repo_dir, check=True) - - subprocess.run( - [ - sys.executable, - "export_trt.py", - "--trt-path", - str(engine_path.absolute()), - "--onnx-path", - str(onnx_path), - ], - cwd=repo_dir, - check=True, - ) - - logger.info(f"Successfully built engine: {engine_path}") - -if __name__ == "__main__": - fire.Fire(compile_depth)