nodes/yuvraj108c-ComfyUI-Rife-Tensorrt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cuda graphs multi stream

Browse Source
This commit is contained in:
yuvraj108c
2024-10-04 13:50:17 +04:00
parent 272d16ab6b 08fa66cd1e
commit 11b6a8fde8
3 changed files with 52 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+8 -11
View File
@@ -6,6 +6,7 @@ from .vfi_utilities import preprocess_frames, postprocess_frames, generate_frame
from .trt_utilities import Engine, MultiStreamEngine
import folder_paths
import time
from polygraphy import cuda
ENGINE_DIR = os.path.join(folder_paths.models_dir, "tensorrt", "rife")
@@ -16,9 +17,9 @@ class RifeTensorrt:
"required": {
"frames": ("IMAGE", ),
"engine": (os.listdir(ENGINE_DIR),),
"clear_cache_after_n_frames": ("INT", {"default": 50, "min": 1, "max": 1000}),
"clear_cache_after_n_frames": ("INT", {"default": 100, "min": 1, "max": 1000}),
"multiplier": ("INT", {"default": 2, "min": 1}),
"cuda_streams": ("INT", {"default": 1, "min": 1, "max": 24}),
"use_cuda_graph": ("BOOLEAN", {"default": True}),
"keep_model_loaded": ("BOOLEAN", {"default": False}),
},
}
@@ -26,16 +27,16 @@ class RifeTensorrt:
RETURN_TYPES = ("IMAGE", )
FUNCTION = "vfi"
CATEGORY = "tensorrt"
OUTPUT_NODE = True
OUTPUT_NODE=True
def vfi(
self,
frames,
engine,
clear_cache_after_n_frames=50,
clear_cache_after_n_frames=100,
multiplier=2,
cuda_streams=1,
keep_model_loaded=False
use_cuda_graph=True,
keep_model_loaded=False,
):
B, H, W, C = frames.shape
shape_dict = {
@@ -44,12 +45,10 @@ class RifeTensorrt:
"output": {"shape": (1, 3, H, W)},
}
# setup tensorrt engine
cudaStream = torch.cuda.current_stream().cuda_stream
cudaStream = cuda.Stream()
engine_path = os.path.join(ENGINE_DIR, engine)
if (not hasattr(self, 'engine') or self.engine_label != engine):
self.engine = MultiStreamEngine(engine_path)
# self.engine = Engine(engine_path)
logger(f"Loading TensorRT engine: {engine_path}")
self.engine.load()
self.engine_label = engine
@@ -65,9 +64,7 @@ class RifeTensorrt:
frames = preprocess_frames(frames)
logger("preprocessing done")
def return_middle_frame(data_batch):
# timestep_t = torch.tensor([timestep], dtype=torch.float32).to(get_torch_device())
# s = time.time()
# output = self.engine.infer({"img0": frame_0, "img1": frame_1, "timestep": timestep_t}, cudaStream)
results = self.engine.infer(data_batch)
# e = time.time()
# print(f"Time taken to infer: {(e-s)*1000} ms")
requirements.txt
+2 -1
View File
@@ -1,4 +1,5 @@
einops
colored
polygraphy
tensorrt==10.4.0
tensorrt==10.4.0
cuda-python
trt_utilities.py
+42 -16
View File
@@ -18,6 +18,8 @@ from tqdm import tqdm
import copy
from collections import OrderedDict
from typing import List, Dict
from cuda import cudart
from polygraphy import cuda
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
G_LOGGER.module_severity = G_LOGGER.ERROR
@@ -45,6 +47,17 @@ torch_to_numpy_dtype_dict = {
value: key for (key, value) in numpy_to_torch_dtype_dict.items()
}
# https://github.com/Jeff-LiangF/streamv2v/blob/18c1a3bd56ff348d54a3300605936980bb13b03c/src/streamv2v/acceleration/tensorrt/utilities.py
def CUASSERT(cuda_ret):
err = cuda_ret[0]
if err != cudart.cudaError_t.cudaSuccess:
raise RuntimeError(
f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t"
)
if len(cuda_ret) > 1:
return cuda_ret[1]
return None
class TQDMProgressMonitor(trt.IProgressMonitor):
def __init__(self):
trt.IProgressMonitor.__init__(self)
@@ -115,7 +128,6 @@ class TQDMProgressMonitor(trt.IProgressMonitor):
# There is no need to propagate this exception to TensorRT. We can simply cancel the build.
return False
class Engine:
def __init__(
self,
@@ -208,7 +220,6 @@ class Engine:
return 0
def load(self):
# logger(f"Loading TensorRT engine: {self.engine_path}")
self.engine = engine_from_bytes(bytes_from_path(self.engine_path))
def activate(self, reuse_device_memory=None):
@@ -238,19 +249,11 @@ class Engine:
nvtx.range_pop()
def infer(self, feed_dict, stream, use_cuda_graph=False):
nvtx.range_push("set_tensors")
for name, buf in feed_dict.items():
self.tensors[name].copy_(buf)
for name, tensor in self.tensors.items():
self.context.set_tensor_address(name, tensor.data_ptr())
nvtx.range_pop()
nvtx.range_push("execute")
noerror = self.context.execute_async_v3(stream)
if not noerror:
raise ValueError("ERROR: inference failed.")
nvtx.range_pop()
return self.tensors
class MultiStreamEngine(Engine):
def __init__(self, engine_path, num_streams=2):
@@ -266,7 +269,8 @@ class MultiStreamEngine(Engine):
def activate(self):
super().activate()
for _ in range(self.num_streams):
stream = torch.cuda.Stream()
# stream = torch.cuda.Stream()
stream = cuda.stream()
context = self.engine.create_execution_context()
self.streams.append(stream)
self.contexts.append(context)
@@ -306,7 +310,7 @@ class MultiStreamEngine(Engine):
tensors[binding] = tensor
nvtx.range_pop()
def infer(self, feed_dicts: List[Dict[str, torch.Tensor]], use_cuda_graph=False):
def infer(self, feed_dicts: List[Dict[str, torch.Tensor]], use_cuda_graph=True):
results = []
num_batches = len(feed_dicts)
@@ -327,10 +331,32 @@ class MultiStreamEngine(Engine):
for name, tensor in tensors.items():
context.set_tensor_address(name, tensor.data_ptr())
if use_cuda_graph:
if self.cuda_graph_instance is not None:
CUASSERT(cudart.cudaGraphLaunch(self.cuda_graph_instance, stream.ptr))
CUASSERT(cudart.cudaStreamSynchronize(stream.ptr))
print("using graph")
else:
# do inference before CUDA graph capture
noerror = self.context.execute_async_v3(stream.ptr)
if not noerror:
raise ValueError("ERROR: inference failed.")
# capture cuda graph
CUASSERT(
cudart.cudaStreamBeginCapture(stream.ptr, cudart.cudaStreamCaptureMode.cudaStreamCaptureModeGlobal)
)
self.context.execute_async_v3(stream.ptr)
self.graph = CUASSERT(cudart.cudaStreamEndCapture(stream.ptr))
self.cuda_graph_instance = CUASSERT(cudart.cudaGraphInstantiate(self.graph, 0))
else:
noerror = self.context.execute_async_v3(stream.ptr)
if not noerror:
raise ValueError("ERROR: inference failed.")
success = context.execute_async_v3(stream.cuda_stream)
if not success:
raise RuntimeError(f"Inference failed for batch {i + j}")
# success = context.execute_async_v3(stream.cuda_stream)
# if not success:
# raise RuntimeError(f"Inference failed for batch {i + j}")
batch_results.append({name: tensor.clone() for name, tensor in tensors.items()
if self.engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT})
@@ -341,4 +367,4 @@ class MultiStreamEngine(Engine):
results.extend(batch_results)
return results
return results