diff --git a/__init__.py b/__init__.py index 234d4c7..b37ac9a 100644 --- a/__init__.py +++ b/__init__.py @@ -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") diff --git a/requirements.txt b/requirements.txt index a0e8f28..b479d99 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ einops colored polygraphy -tensorrt==10.4.0 \ No newline at end of file +tensorrt==10.4.0 +cuda-python \ No newline at end of file diff --git a/trt_utilities.py b/trt_utilities.py index e8222b2..2258d08 100644 --- a/trt_utilities.py +++ b/trt_utilities.py @@ -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 \ No newline at end of file + return results